perf(mode): pin sliced mode-solver plane to single-device sharding

src/fdtdx/core/jax/sharding.py

@@ -41,6 +41,60 @@ def get_named_sharding_from_shape(
 counter = 0
 
 
+def _replicated_mesh_sharding(compute_devices: list[jax.Device]) -> jax.sharding.NamedSharding:
+    """Build a REPLICATED NamedSharding over ``compute_devices`` -- NOT SingleDeviceSharding, whose
+    ``{device 0}`` device set is incompatible with the ``[0,1,...]`` set of the surrounding x-sharded
+    computation (jit raises "Received incompatible devices for jitted computation"). A replicated
+    PartitionSpec() keeps data on the full mesh but UN-sharded, so a subsequent placement (see
+    ``pin_to_single_device``) gathers only the small operand, never the full x-sharded domain, while
+    staying device-set-compatible with the parent computation.
+
+    Split out of ``pin_to_single_device`` so the (cheap, pure) mesh/sharding construction can be unit-
+    tested directly with a MOCKED device list (including a duplicated physical device, which the actual
+    data placement in ``pin_to_single_device`` cannot accept -- see ``jax.device_put``'s "distinct
+    devices" requirement, exercised only in the subprocess-based multi-device tests).
+    """
+    devices = mesh_utils.create_device_mesh((len(compute_devices),), devices=compute_devices)
+    mesh = jax.sharding.Mesh(devices=devices, axis_names=(SHARD_STR,))
+    return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+
+
+def pin_to_single_device(arr: jax.Array) -> jax.Array:
+    """Pin ``arr`` to a single-device replicated layout via ``jax.device_put``.
+
+    On a multi-device run the field/material arrays are x-sharded (NamedSharding
+    ``P(None,'shard',None,None)``). When a thin transverse plane is sliced out of such an
+    array and fed to a host ``pure_callback`` (the Tidy3D mode solver, which runs at
+    ``{maximal device=0}``), XLA must materialise a single-device operand for the callback.
+    Without an explicit constraint XLA may conservatively all-gather the *parent* array to
+    satisfy the callback's single-device requirement. Pinning the already-sliced PLANE to a
+    single-device replicated sharding tells XLA to gather only the thin band (a few-hundred-KB
+    plane), never the full ~100M-cell domain.
+
+    Also used on the pure_callback's OWN outputs (see ``compute_mode``): those carry a maximal/
+    single-device GSPMDSharding that ``jax.lax.with_sharding_constraint`` cannot convert (raises
+    "Cannot convert GSPMDSharding {maximal device=0} into SdyArray" via XLA's Shardy partitioner).
+    ``jax.device_put`` performs the placement/reshard directly and accepts both this maximal-sharded
+    case and the ordinary already-NamedSharding case (the sliced input this function was originally
+    written for), so one implementation covers both call patterns.
+
+    No-op on a single-device run (``len(jax.devices()) == 1``): returns ``arr`` unchanged so the
+    single-GPU path stays byte-identical (no sharding constraint is inserted at all).
+    """
+    if len(jax.devices()) == 1:
+        return arr
+    replicated = _replicated_mesh_sharding(jax.devices())
+    # device_put, NOT jax.lax.with_sharding_constraint: the latter internally converts through XLA's
+    # Shardy partitioner (`api.jit(_identity_fn, out_shardings=sharding)`), which raises "Cannot convert
+    # GSPMDSharding {maximal device=0} into SdyArray" when `arr` is the RAW OUTPUT of a jax.pure_callback
+    # (pure_callback results carry a maximal/single-device GSPMDSharding, not an ordinary NamedSharding --
+    # a fundamentally different representation Shardy's conversion path does not handle). device_put
+    # performs the actual data placement/reshard directly and accepts a maximal-sharded source array,
+    # so it works for BOTH the already-well-formed NamedSharding case (the sliced x-sharded input this
+    # function was originally written for) and the pure_callback-output case (its raw mode-solver result).
+    return jax.device_put(arr, replicated)
+
+
 def get_dtype_bytes(dtype: jnp.dtype) -> int:
     return jnp.dtype(dtype).itemsize
 

src/fdtdx/core/physics/modes.py

@@ -10,6 +10,7 @@
 from tidy3d.components.mode.solver import compute_modes as _compute_modes
 
 from fdtdx.core.axis import get_transverse_axes
+from fdtdx.core.jax.sharding import pin_to_single_device
 from fdtdx.core.misc import expand_to_3x3
 from fdtdx.core.physics.metrics import normalize_by_poynting_flux
 
@@ -361,6 +362,14 @@ def mode_helper(permittivity, permeability, c0_um, c1_um):
         jax.lax.stop_gradient(c0_um),
         jax.lax.stop_gradient(c1_um),
     )
+    # pure_callback executes on a single host device ({maximal device=0}); its OUTPUT, like the sliced
+    # INPUT above, needs an explicit sharding reconciliation before further ops run it through the
+    # surrounding x-sharded multi-device computation -- otherwise XLA/Shardy cannot convert the callback's
+    # single-device GSPMDSharding into an SdyArray at the next op (observed: jnp.expand_dims below raises
+    # "Cannot convert GSPMDSharding {maximal device=0} into SdyArray" on >1 device).
+    mode_E_raw = pin_to_single_device(mode_E_raw)
+    mode_H_raw = pin_to_single_device(mode_H_raw)
+    eff_idx = pin_to_single_device(eff_idx)   # third callback output -- same fix, easy to miss
     mode_E = jnp.expand_dims(mode_E_raw, axis=propagation_axis + 1)
     mode_H = jnp.expand_dims(mode_H_raw, axis=propagation_axis + 1)
 

src/fdtdx/objects/detectors/mode.py

@@ -7,6 +7,7 @@
 from fdtdx import constants
 from fdtdx.config import SimulationConfig
 from fdtdx.core.jax.pytrees import autoinit, frozen_field, private_field
+from fdtdx.core.jax.sharding import pin_to_single_device
 from fdtdx.core.null import Null
 from fdtdx.core.physics.modes import compute_mode
 from fdtdx.dispersion import effective_complex_inv_permittivity
@@ -156,9 +157,10 @@ def apply(
         dispersive_c4: jax.Array | None = None,
     ) -> Self:
         del key
-        inv_permittivity_slice = inv_permittivities[:, *self.grid_slice]
+        # Pin the sliced plane to a single device (see pin_to_single_device docstring for why).
+        inv_permittivity_slice = pin_to_single_device(inv_permittivities[:, *self.grid_slice])
         if isinstance(inv_permeabilities, jax.Array) and inv_permeabilities.ndim > 0:
-            inv_permeability_slice = inv_permeabilities[:, *self.grid_slice]
+            inv_permeability_slice = pin_to_single_device(inv_permeabilities[:, *self.grid_slice])
         else:
             inv_permeability_slice = inv_permeabilities
 

src/fdtdx/objects/sources/mode.py

@@ -10,6 +10,7 @@
 from fdtdx.core.axis import get_transverse_axes
 from fdtdx.core.grid import calculate_time_offset_yee
 from fdtdx.core.jax.pytrees import autoinit, frozen_field, private_field
+from fdtdx.core.jax.sharding import pin_to_single_device
 from fdtdx.core.linalg import get_wave_vector_raw
 from fdtdx.core.physics.metrics import compute_energy
 from fdtdx.core.physics.modes import compute_mode
@@ -112,10 +113,11 @@ def apply(
             raise NotImplementedError()
 
         # inv_permittivities shape: (3, Nx, Ny, Nz) - slice with component dimension
-        inv_permittivity_slice = inv_permittivities[:, *self.grid_slice]
+        # Pin the sliced plane to a single device (see pin_to_single_device docstring for why).
+        inv_permittivity_slice = pin_to_single_device(inv_permittivities[:, *self.grid_slice])
         if isinstance(inv_permeabilities, jax.Array) and inv_permeabilities.ndim > 0:
             # inv_permeabilities shape: (3, Nx, Ny, Nz) - slice with component dimension
-            inv_permeability_slice = inv_permeabilities[:, *self.grid_slice]
+            inv_permeability_slice = pin_to_single_device(inv_permeabilities[:, *self.grid_slice])
         else:
             inv_permeability_slice = inv_permeabilities
 

tests/unit/core/jax/test_sharding.py

@@ -1,5 +1,9 @@
 """Unit tests for fdtdx.core.jax.sharding module."""
 
+import os
+import subprocess
+import sys
+import textwrap
 from unittest import mock
 from unittest.mock import MagicMock
 
@@ -10,9 +14,11 @@
 import fdtdx.core.jax.sharding as sharding_module
 from fdtdx.constants import SHARD_STR
 from fdtdx.core.jax.sharding import (
+    _replicated_mesh_sharding,
     create_named_sharded_matrix,
     get_dtype_bytes,
     get_named_sharding_from_shape,
+    pin_to_single_device,
     pretty_print_sharding,
 )
 
@@ -198,3 +204,184 @@ def test_sharding_axis_fallback_when_dim_is_one(self):
     def test_has_named_sharding(self):
         result = create_named_sharded_matrix(shape=(4, 6), value=1.0, sharding_axis=0, dtype=jnp.float32, backend="cpu")
         assert isinstance(result.sharding, jax.sharding.NamedSharding)
+
+
+# ---- _replicated_mesh_sharding ----
+
+
+class TestReplicatedMeshSharding:
+    """Tests for the _replicated_mesh_sharding helper, split out of pin_to_single_device so this
+    (pure, cheap) mesh/sharding construction can be covered in-process -- unlike pin_to_single_device's
+    own data placement, it does not require genuinely distinct physical devices, so a MOCKED duplicated
+    device list is a faithful, subprocess-free test (the corresponding data placement is exercised
+    separately by TestPinToSingleDevice's subprocess-based multi-device tests below)."""
+
+    def test_single_device_replicated_spec(self):
+        cpu = CPU_DEVICES[0]
+        result = _replicated_mesh_sharding([cpu])
+        assert isinstance(result, jax.sharding.NamedSharding)
+        assert tuple(result.spec) == ()
+        assert result.mesh.devices.shape == (1,)
+
+    def test_duplicated_device_list_replicated_spec(self):
+        # A duplicated physical device (simulating >1 "device" without needing real distinct hardware)
+        # is accepted by mesh/sharding CONSTRUCTION -- only the actual data placement in
+        # pin_to_single_device requires genuine distinctness (see module docstring above).
+        cpu = CPU_DEVICES[0]
+        result = _replicated_mesh_sharding([cpu, cpu])
+        assert isinstance(result, jax.sharding.NamedSharding)
+        assert tuple(result.spec) == ()
+        assert result.mesh.devices.shape == (2,)
+        assert SHARD_STR in result.mesh.axis_names
+
+    def test_result_usable_as_pin_to_single_device_input(self):
+        """pin_to_single_device calls this helper directly with jax.devices() -- confirm the two agree
+        on shape/axis-name conventions for a real (single, non-mocked) device list."""
+        expected = _replicated_mesh_sharding(jax.devices())
+        assert isinstance(expected, jax.sharding.NamedSharding)
+        assert expected.mesh.devices.shape == (len(jax.devices()),)
+
+
+# ---- pin_to_single_device ----
+
+
+def _run_in_subprocess_with_devices(num_devices: int, body: str) -> None:
+    """Run ``body`` in a fresh interpreter with ``num_devices`` virtual CPU devices.
+
+    ``jax.devices()`` reflects the platform's device count only at JAX's first backend
+    initialisation, which has already happened (with a single forced CPU device, see
+    ``tests/conftest.py``) by the time this test module runs. A real multi-device mesh can
+    therefore only be exercised in a separate process started with
+    ``XLA_FLAGS=--xla_force_host_platform_device_count=N`` set beforehand. Mocking
+    ``jax.devices()`` to return the same physical device twice is not a substitute here:
+    unlike the mesh-shape/spec checks elsewhere in this file, ``pin_to_single_device`` places
+    data onto a real device mesh eagerly, which requires genuinely distinct physical devices
+    (a mesh built from a duplicated device raises at construction/placement time).
+    """
+    env = {
+        **os.environ,
+        "JAX_PLATFORMS": "cpu",
+        "XLA_FLAGS": f"--xla_force_host_platform_device_count={num_devices}",
+    }
+    result = subprocess.run(
+        [sys.executable, "-c", textwrap.dedent(body)],
+        capture_output=True,
+        text=True,
+        env=env,
+        timeout=60,
+    )
+    assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+
+
+class TestPinToSingleDevice:
+    """Tests for the pin_to_single_device helper."""
+
+    def test_single_device_is_noop_identity(self):
+        # The default test session forces a single CPU device (see tests/conftest.py).
+        assert len(jax.devices()) == 1
+        arr = jnp.arange(12.0).reshape(3, 4)
+        result = pin_to_single_device(arr)
+        assert result is arr
+
+    def test_multi_device_preserves_shape_dtype_and_values(self):
+        _run_in_subprocess_with_devices(
+            2,
+            """
+            import jax, jax.numpy as jnp
+            from fdtdx.core.jax.sharding import pin_to_single_device
+
+            assert len(jax.devices()) == 2
+            arr = jnp.arange(24.0, dtype=jnp.float32).reshape(4, 6)
+            out = pin_to_single_device(arr)
+            assert out.shape == arr.shape
+            assert out.dtype == arr.dtype
+            assert bool(jnp.array_equal(out, arr))
+            """,
+        )
+
+    def test_multi_device_result_is_replicated_named_sharding(self):
+        _run_in_subprocess_with_devices(
+            2,
+            """
+            import jax, jax.numpy as jnp
+            from fdtdx.core.jax.sharding import pin_to_single_device
+
+            out = pin_to_single_device(jnp.zeros((4, 6), dtype=jnp.float32))
+            assert isinstance(out.sharding, jax.sharding.NamedSharding)
+            # Fully replicated: an empty PartitionSpec, not sharded along any axis.
+            assert tuple(out.sharding.spec) == ()
+            assert "shard" in out.sharding.mesh.axis_names
+            """,
+        )
+
+    def test_multi_device_compatible_with_sliced_x_sharded_parent(self):
+        """Regression test for the exact usage in ModePlaneSource / ModeOverlapDetector:
+        pinning a plane sliced out of an x-sharded array inside a jit must not raise JAX's
+        "Received incompatible devices for jitted computation" error, which is what happens
+        if the sliced plane were instead placed on a plain SingleDeviceSharding (device 0)
+        while the surrounding computation is compiled over the full device mesh.
+        """
+        _run_in_subprocess_with_devices(
+            2,
+            """
+            import jax, jax.numpy as jnp
+            from fdtdx.core.jax.sharding import create_named_sharded_matrix, pin_to_single_device
+
+            # (component, Nx, Ny, Nz) permittivity-like array, x-sharded on axis=1.
+            arr = create_named_sharded_matrix(
+                (3, 8, 6, 6), value=2.0, dtype=jnp.float32, sharding_axis=1, backend="cpu"
+            )
+
+            @jax.jit
+            def slice_and_pin(a):
+                plane = a[:, 2:3, :, :]
+                return pin_to_single_device(plane)
+
+            out = slice_and_pin(arr)
+            assert out.shape == (3, 1, 6, 6)
+            assert bool(jnp.allclose(out, 2.0))
+            """,
+        )
+
+    def test_multi_device_compatible_with_pure_callback_output(self):
+        """Regression test for compute_mode's usage: pinning the RAW OUTPUT of a
+        jax.pure_callback (as returned by the Tidy3D mode solver) must not raise "Cannot
+        convert GSPMDSharding {maximal device=0} into SdyArray". A pure_callback's result
+        carries a maximal/single-device GSPMDSharding -- a fundamentally different
+        representation than the ordinary NamedSharding case covered by the sliced-input
+        test above, which XLA's Shardy partitioner cannot convert via
+        jax.lax.with_sharding_constraint (this failed before this fix; device_put succeeds
+        because it performs the placement directly instead of a Shardy-mediated re-shard).
+        """
+        _run_in_subprocess_with_devices(
+            2,
+            """
+            import numpy as np
+            import jax, jax.numpy as jnp
+            from fdtdx.core.jax.sharding import create_named_sharded_matrix, pin_to_single_device
+
+            arr = create_named_sharded_matrix(
+                (3, 8, 6, 6), value=2.0, dtype=jnp.float32, sharding_axis=1, backend="cpu"
+            )
+
+            def host_fn(x):
+                # Stand-in for the Tidy3D mode-solver call in compute_mode: an arbitrary
+                # host-side computation returning a NEW array (not just a passthrough).
+                return np.asarray(x) * 3.0
+
+            @jax.jit
+            def slice_callback_and_pin(a):
+                plane = a[:, 2:3, :, :]
+                raw = jax.pure_callback(
+                    host_fn, jax.ShapeDtypeStruct(plane.shape, plane.dtype), plane
+                )
+                pinned = pin_to_single_device(raw)
+                # Mirrors compute_mode's jnp.expand_dims(mode_E_raw, ...) right after pinning --
+                # the exact next op that raised before this fix.
+                return jnp.expand_dims(pinned, axis=0)
+
+            out = slice_callback_and_pin(arr)
+            assert out.shape == (1, 3, 1, 6, 6)
+            assert bool(jnp.allclose(out, 6.0))
+            """,
+        )