[rogeriojorge]

What changed

• shard DG distribution and field state over a selected spatial cell axis (x, y, or z)
• exchange the one-cell periodic DG halo explicitly at device boundaries
• use one global RMS error norm so adaptive steps remain synchronized
• stage replicated coefficient arrays from the host before entering a cached shard_map solve
• add one small CPU/GPU strong-scaling example and one serial-equivalence test

This is stacked on #39 so the sharding diff can be reviewed independently. The default serial behavior is unchanged.

Root cause

Two details made the earlier implementation invalid on multiple GPUs. A local roll wrapped at each device boundary instead of the periodic domain boundary, and placing an already materialized single-GPU JAX array with replicated NamedSharding left zero-valued replicas on GPU 1 in the office environment. Diffrax still reported success while the 2D state became NaN. The revised path stages arrays from host memory, declares layouts with an Explicit mesh, and exchanges only the boundary cell with ppermute.

The sharded solve is a module-level JIT so repeated calls reuse compilation; constructing shard_map inside simulation caused repeated compilation and hid scaling behind JIT overhead.

Validation

• pytest -q: 5 passed
• focused regression under JAX_NUM_CPU_DEVICES=2: passed
• 2D 8x8 H4 two-mode DG GEM fixed step on two CPUs: bitwise identical to serial; both fields and distributions finite
• 2D adaptive GEM smoke run on two CPUs: bitwise identical to serial
• local CPU benchmark, warm: 0.448 s (1 device), 0.269 s (4), 0.226 s (8); 1.98x speedup at 8
• office RTX A4000 benchmark, warm: 0.648 s (1 GPU), 0.344 s (2); 1.88x speedup
• 2D 8x8 H4 two-mode DG GEM step on two GPUs: finite on both shards; 0.074 s warm after a 10.3 s cold compile
• output buffers verified at 1/4 and 1/8 local spatial extents on CPUs and 1/2 extents on GPUs

Scope

This PR establishes correct single-axis spatial decomposition. Larger nonlinear GEM production runs and physics validation remain in the separate example PR so this numerical infrastructure change stays reviewable.

[rogeriojorge]

Squashed into #39, which now contains the sharding change.