[pallas] Preserve out_avals memory spaces in `_mpmd_map_discharge_rule` #40783

Open
copybara-service[bot] opened 7:18pm on September 17, 2026 wants to merge 66 Δ into google/jax main from
test_983558282

Pull Request #40783 commits