[Benjamin-Herfray]

Part 1 of 7 of multi-device (spatial + species) execution for the DG solver. The PRs are stacked; merge them in order. This one targets Discontinuous-Galerkin.

The Hermite streaming matrices A_x, A_y, A_z and the Maxwell flux matrices R are real symmetric. They are now split into positive- and negative-eigenvalue parts with jnp.linalg.eigh (new helper split_flux_matrix) instead of a general eig. eigh gives orthonormal eigenvectors and runs on CPU and GPU backends. The central/upwind choices are unchanged.

Tests: tests/test_parallel.py checks symmetry, A_p + A_m = A, sign definiteness and agreement with the previous eig construction. The full test suite passes (7 tests).

[rogeriojorge]

Profiled the #74-#80 stack with 1, 2 and 4 emulated devices; details on #78. This one looks fine: eigh is the right choice (orthonormal, runs on GPU) and the flux split is exact.

Needed: a test that differentiates through compute_A_pm_matrices with respect to alpha_s and u_s. eigh gradients are only safe while the eigenvalues stay distinct, so this should be pinned.

[rogeriojorge]

@Benjamin-Herfray, handing the DG work over to you. Here is the state as of 2026-10-26, from merging each PR into its base and running the suite on jax 0.12.0.1.

The chain: #74 → #75 → #76 → #77 → #78 → #79 → #80, then #87 on top, which has the jax.grad and jit fixes from the review on #78/#80.

1. Three tests fail from Integrate one sharded (Ck, Fk) PyTree solver for one or many devices #78 onward on jax 0.12. These are the three cases (2x1x1x1, 2x1x2x1, 2x2x1x1) of tests/test_parallel.py::test_compiled_rhs_communicates_only_faces_and_species_currents. The failing check is the collective-permute participant check sorted(x[0] for x in cs) == list(range(ps)): newer XLA emits a different source_target_pairs pattern. The physics and the sharded-versus-single-device results are unaffected. I suggest checking the communication pattern by type and size only, without exact participant lists, or pinning the check to jax versions where it holds.
2. The base is stale. Discontinuous-Galerkin is 50 commits behind main. Sync Discontinuous-Galerkin with main #54, a sync of main into the DG branch, is not in the chain and now conflicts. Suggested order: refresh Sync Discontinuous-Galerkin with main #54 (merge main into it), merge it into Discontinuous-Galerkin, then merge the base into each chain PR in order. The most recent changes on main touch pyproject.toml (Python 3.10, dependencies), _simulation.py (dtmin, max_steps, throw, the (Ck, Fk) state) and the tests.
3. Closed as superseded: Keep DG state as a PyTree and shard it over spatial cells #39 (single-axis sharding, replaced by Integrate one sharded (Ck, Fk) PyTree solver for one or many devices #78) and Differentiable DG simulation: solver keywords and eigh upwind matrices #45 (DG autodiff, replaced by Differentiable multi-device DG simulation, Python 3.10 #87).
4. Still unmeasured: real multi-GPU scaling. Only emulated CPU devices were profiled.

The merge order within the DG branch is yours to decide. I won't push to these branches unless you ask.