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

The solver adaptive-cutoff path seeded its fori_loop bracket from
constants (jnp.zeros / jnp.full), while 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 disagreed on varying manual axes and scan's
carry type check rejected the loop at trace time:

  TypeError: scan body function carry input and carry output must have
  equal types [...] float32[1024] vs float32[1024]{V:dp}

Any data-parallel run using adaptive_cutoff_method="solver" failed to
trace; single-device runs were unaffected, since callers only apply
shard_map when there is more than one device.

pcast the bracket to varying, reading the axes off the data so the fix
is agnostic to the caller's mesh axis names and a no-op outside
shard_map. Runtime behaviour is unchanged: each shard already solves on
its own atoms, and the annotation carries no computation.

Adds tests/test_shard_map.py, which runs both adaptive-cutoff methods
under shard_map on forced CPU devices and checks values and gradients
against the single-device result -- multi-device coverage the suite
lacked. Verified on GPU: the trace-time failure is gone on a 4-way
sharded PET training run.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: E-Rum

src/petjax/select.py

@@ -298,9 +298,14 @@ def get_adaptive_cutoffs_solver(
     r_ij_d = jax.lax.stop_gradient(r_ij)
 
     # Bracket [r_lo, r_hi] with f(r_lo) <= 0 <= f(r_hi): n_total(0) = 0 and the
-    # baseline alone reaches num_neighbors at r = cutoff.
-    r_lo = jnp.zeros(num_atoms, dtype=r_ij.dtype)
-    r_hi = jnp.full(num_atoms, cutoff, dtype=r_ij.dtype)
+    # baseline alone reaches num_neighbors at r = cutoff. The carry must enter
+    # varying to match the body's outputs (derived from this shard's r_ij);
+    # axes read off the data, so no-op outside shard_map.
+    varying = tuple(jax.typeof(r_ij_d).vma)
+    r_lo = jax.lax.pcast(jnp.zeros(num_atoms, dtype=r_ij.dtype), varying, to="varying")
+    r_hi = jax.lax.pcast(
+        jnp.full(num_atoms, cutoff, dtype=r_ij.dtype), varying, to="varying"
+    )
 
     # 10 iterations converge to fp32 precision (upstream's choice). Newton
     # steps that would leave the bracket (flat shoulders between bumps) fall

tests/test_shard_map.py

@@ -0,0 +1,85 @@
+"""Adaptive cutoffs under shard_map (data-parallel training).
+
+Callers such as iris wrap the step in `jax.shard_map` over a data-parallel
+mesh axis, which makes every array derived from the local batch "varying"
+over that axis. Code traced inside must keep loop carries varying too, so
+these run on CPU with a forced device count -- no GPU needed.
+"""
+
+import os
+
+os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=4")
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from jax.sharding import NamedSharding, PartitionSpec as P
+
+from petjax.select import get_adaptive_cutoffs, get_adaptive_cutoffs_solver
+
+N_ATOMS = 64
+N_PAIRS = 512
+CUTOFF = 7.5
+CUTOFF_WIDTH = 1.0
+NUM_NEIGHBORS = 16
+
+pytestmark = pytest.mark.skipif(
+    jax.device_count() < 2, reason="needs >= 2 devices (XLA_FLAGS device count)"
+)
+
+
+def _inputs(n_dev):
+    rng = np.random.default_rng(0)
+    return (
+        jnp.asarray(rng.integers(0, N_ATOMS, size=(n_dev, N_PAIRS)), dtype=jnp.int32),
+        jnp.asarray(rng.uniform(0.5, CUTOFF, size=(n_dev, N_PAIRS)), dtype=jnp.float32),
+        jnp.asarray(rng.random((n_dev, N_PAIRS)) > 0.1),
+    )
+
+
+def _sharded(fn, mesh):
+    """Mirror the training wrapper: shard on the leading axis, squeeze, call."""
+
+    @jax.jit
+    def wrapped(*args):
+        @jax.shard_map(
+            mesh=mesh, in_specs=(P("dp"),) * len(args), out_specs=P("dp")
+        )
+        def inner(*args):
+            return fn(*[jnp.squeeze(a, 0) for a in args])[None]
+
+        return inner(*args)
+
+    return wrapped
+
+
+@pytest.mark.parametrize(
+    "method",
+    [get_adaptive_cutoffs_solver, get_adaptive_cutoffs],
+    ids=["solver", "grid"],
+)
+def test_adaptive_cutoffs_match_single_device(method):
+    n_dev = jax.device_count()
+    centers, r_ij, pair_mask = _inputs(n_dev)
+
+    def solve(centers, r_ij, pair_mask):
+        return method(
+            centers, r_ij, pair_mask, NUM_NEIGHBORS, N_ATOMS, CUTOFF, CUTOFF_WIDTH
+        )
+
+    ref = jnp.stack([solve(centers[i], r_ij[i], pair_mask[i]) for i in range(n_dev)])
+
+    mesh = jax.make_mesh((n_dev,), ("dp",))
+    shard = NamedSharding(mesh, P("dp"))
+    args = [jax.device_put(x, shard) for x in (centers, r_ij, pair_mask)]
+    got = _sharded(solve, mesh)(*args)
+
+    # 1 ULP: XLA fuses the segment sums differently under shard_map
+    np.testing.assert_allclose(np.asarray(got), np.asarray(ref), rtol=1e-6, atol=1e-6)
+
+    with jax.set_mesh(mesh):  # the scalar output lives on the mesh
+        grad = jax.jit(
+            jax.grad(lambda r: _sharded(solve, mesh)(args[0], r, args[2]).sum())
+        )(args[1])
+    assert jnp.all(jnp.isfinite(grad))