[Pallas][MPMD] Preserve out_avals memory spaces in _mpmd_map_discharge_rule #40777

Open
copybara-service[bot] opened 4:27pm on September 17, 2026 wants to merge 1 commit into google/jax main from
test_983480985

Pull Request #40777 commits