[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 quality checks

Code files and functions touched by this pull request are analyzed below.

Quality targets

How does this pull request compare to chosen code quality targets? Read more about code quality targets.

Target name Target value This PR
Maximum Diff Delta per PR 500 66
Maximum comment count per PR 10 0
Pull request maximum business days between activity 3 0
Pull request maximum business days open 5 0

Coupled files

Are there any files from this commit that are always committed alongside another file? This often suggests that there exists code coupling to potentially remove. Read more about code coupling

No file coupling detected. 🎉

Biggest files changes in pull request

Path View whole file Lines removed Lines added Final line count
jax/_src/pallas/mpmd.py 5 9 ---
tests/pallas/tpu_pallas_state_test.py 0 29 ---
tests/pallas/tpu_pallas_mpmd_test.py 0 7 ---

️ Biggest functions grown or shrunk by pull request

Which functions had their line count change as part of this pull request?

Could not find any functions changed by this pull request.