select: keep the solver's Newton carry varying under shard_map

[E-Rum]

Problem

get_adaptive_cutoffs_solver seeds its fori_loop bracket from constants:

r_lo = jnp.zeros(num_atoms, dtype=r_ij.dtype)
r_hi = jnp.full(num_atoms, cutoff, dtype=r_ij.dtype)

but the loop body's outputs derive from r_ij. Under shard_map, r_ij varies over the data-parallel mesh axis, so carry-in and carry-out disagree on manual axis types and scan's carry type check rejects the loop at trace time:

TypeError: scan body function carry input and carry output must have equal types, but they differ:
  * the input carry component loop_carry[1][0] has type float32[1024] but the corresponding
    output carry component has type float32[1024]{V:dp}

Every data-parallel run with adaptive_cutoff_method="solver" fails to trace. Single-device runs are unaffected, because callers (e.g. iris) only wrap the step in shard_map when there is more than one device — so this does not show up in a single-GPU smoketest. The grid method has no loop and was never affected.

Fix

pcast the bracket to varying, reading the axes off the data rather than naming them, so it is agnostic to the caller's mesh axis names and a no-op (empty axes) outside shard_map. The third carry element, 0.5 * r_hi, inherits varying-ness from r_hi.

Runtime behaviour is unchanged: each shard already solves on its own atoms, and the annotation carries no computation.

Note this reads jax.typeof(x).manual_axis_type.varying, which is the JAX >= 0.11 spelling (it was .vma in 0.9/0.10). pyproject.toml currently floors jax at 0.4.30; if older versions are still meant to be supported, the floor should be raised rather than this made conditional — on those releases the attribute is absent and the bug is unfixable in the same way.

Tests

tests/test_shard_map.py runs both adaptive-cutoff methods under shard_map and checks values and gradients against the single-device result — multi-device coverage the suite currently lacks. It needs no GPU: conftest.py forces four CPU devices via XLA_FLAGS, set there because XLA reads it at backend init, before any test module imports jax.

• Without the fix: [solver] fails with the TypeError above, [grid] passes.
• With the fix: both pass; full suite 12 passed, 36 skipped on jax 0.11.0.
• Values match single-device to 1 ULP (max 4.8e-07 abs); the drift is XLA fusing the segment sums differently under shard_map.

Verified on hardware (jax 0.11.0, RTX PRO 6000 Blackwell): a 4-way sharded PET training run that previously died at trace time now compiles and proceeds past optimizer/loss setup. A real multi-GPU job is queued to confirm the collectives path end-to-end — the MIG slices used for the quick test cannot form an NCCL communicator, which is a MIG limitation unrelated to this change.

🤖 Generated with Claude Code