Add JVP-safe fixed-index scatter

Author: r-fedorov

README.md

@@ -114,25 +114,31 @@ JVP is useful for more specialized atomistic workflows, including:
 - debugging equivariance by differentiating along an infinitesimal rotation.
 
 MLX 0.31 cannot currently apply JVP directly to a `CustomKernel` primitive.
-For spherical harmonics and tensor products, select the fully differentiable
-MLX implementation:
+For these workflows, select the fully differentiable MLX implementation:
 
 ```python
 import e3nn_mlx
 
-# Spherical harmonics
+# Spherical harmonics and fixed-topology graph reduction
 y = e3nn_mlx.spherical_harmonics(
     degrees, vectors, use_custom_kernel=False
 )
+summed = e3nn_mlx.scatter_sum(
+    messages,
+    edge_dst,
+    num_nodes,
+    use_custom_kernel=False,
+    jvp_safe=True,
+)
 
 # TensorProduct: use this callable inside mx.jvp
 y = tensor_product.differentiable_arrays(left, right, weights)
 ```
 
-MLX 0.31's indexed-add primitive does not implement JVP, so neither scatter
-execution path is a forward-mode fallback. `scatter_sum` supports ordinary
-reverse-mode gradients; compute a message JVP before aggregation or use a
-problem-specific fixed incidence matrix when a graph-reduction JVP is required.
+Because MLX 0.31's indexed-add primitive does not implement JVP, the scatter
+fallback requires eager, fixed indices. It uses a sparse sorted prefix sum with
+linear memory rather than a dense node-by-edge incidence matrix. Ordinary
+scatter calls keep the faster indexed-add implementation.
 
 This changes execution strategy, not mathematical conventions or accuracy.
 Reverse-mode gradients and reverse-over-reverse second derivatives remain

docs/guide/performance.md

@@ -40,10 +40,11 @@ harmonics = o3.spherical_harmonics(
 )
 ```
 
-Scatter is a separate boundary: MLX 0.31's indexed-add primitive does not
-implement JVP, so `scatter_sum` supports reverse-mode differentiation but has
-no forward-mode fallback. Compute a message JVP before aggregation or use a
-problem-specific fixed incidence matrix when the reduction itself needs JVP.
+For graph reduction, use
+`scatter_sum(..., use_custom_kernel=False, jvp_safe=True)`. MLX 0.31's
+indexed-add primitive does not implement JVP, so this fixed-index fallback uses
+a sparse sorted prefix sum with linear memory. Ordinary scatter calls retain
+the faster indexed-add implementation.
 
 See the repository's
 [kernel evaluation guide](https://github.com/lamalab-org/e3nn_mlx/blob/main/evals/KERNEL_EVALUATION.md)

e3nn_mlx/graph.py

@@ -11,8 +11,15 @@ def scatter_sum(
     dim_size: int | None = None,
     *,
     use_custom_kernel: bool = False,
+    jvp_safe: bool = False,
 ):
-    """Sum rows of ``source`` into output rows selected by one-dimensional ``index``."""
+    """Sum rows of ``source`` into rows selected by one-dimensional ``index``.
+
+    Set ``jvp_safe=True`` for forward-mode differentiation with respect to
+    ``source``.  This path requires eager, fixed indices and implements the
+    same reduction with a sorted prefix sum because MLX 0.31's indexed-add
+    primitive has no JVP rule.
+    """
 
     mx, _ = require_mlx()
     if source.ndim < 1:
@@ -25,6 +32,10 @@ def scatter_sum(
         dim_size = 0 if index.size == 0 else int(mx.max(index).item()) + 1
     if dim_size < 0:
         raise ValueError("dim_size must be non-negative")
+    if jvp_safe:
+        if use_custom_kernel:
+            raise ValueError("jvp_safe=True requires use_custom_kernel=False")
+        return _fixed_index_scatter_sum(source, index, dim_size)
     use_metal = (
         use_custom_kernel
         and mlx_metal_available()
@@ -51,6 +62,44 @@ def scatter_sum(
     return output.at[index].add(source)
 
 
+def _fixed_index_scatter_sum(source, index, dim_size: int):
+    """Sparse JVP-safe scatter for an eager, fixed index array."""
+
+    mx, _ = require_mlx()
+    import numpy as np
+
+    try:
+        host_index = np.asarray(index)
+    except Exception as error:
+        raise RuntimeError(
+            "jvp_safe scatter requires an eager fixed index array"
+        ) from error
+    if host_index.size:
+        minimum = int(host_index.min())
+        maximum = int(host_index.max())
+        if minimum < 0 or maximum >= dim_size:
+            raise ValueError("index values must satisfy 0 <= index < dim_size")
+
+    order_host = np.argsort(host_index, kind="stable")
+    sorted_index = host_index[order_host]
+    rows = np.arange(dim_size, dtype=host_index.dtype)
+    starts_host = np.searchsorted(sorted_index, rows, side="left")
+    ends_host = np.searchsorted(sorted_index, rows, side="right")
+    order = mx.array(order_host, dtype=mx.int32)
+    starts = mx.array(starts_host, dtype=mx.int32)
+    ends = mx.array(ends_host, dtype=mx.int32)
+
+    sorted_source = source[order]
+    prefix = mx.concatenate(
+        (
+            mx.zeros((1, *source.shape[1:]), dtype=source.dtype),
+            mx.cumsum(sorted_source, axis=0),
+        ),
+        axis=0,
+    )
+    return prefix[ends] - prefix[starts]
+
+
 def radius_graph(positions, radius: float, batch=None):
     """Build the directed, loop-free batched radius graph used by gate-points models.
 

evals/KERNEL_EVALUATION.md

@@ -27,9 +27,10 @@ derivatives. MLX 0.31 cannot currently apply JVP directly to a `CustomKernel`
 primitive even when its containing custom function supplies a JVP rule. Code
 which needs forward-mode transformation should therefore call tensor-product
 `differentiable_arrays`, or pass `use_custom_kernel=False` to spherical
-harmonics. MLX's indexed-add primitive also lacks JVP, so scatter has no general
-forward-mode fallback in MLX 0.31. This is an explicit runtime boundary, not a
-numerical approximation.
+harmonics. For fixed-topology scatter, pass `use_custom_kernel=False` and
+`jvp_safe=True`; this avoids MLX's missing indexed-add JVP with a sparse sorted
+prefix sum. This is an explicit execution boundary, not a numerical
+approximation.
 
 Typical JVP users are phonon and vibrational-response calculations,
 Hessian-vector products, mixed position/parameter response, tangent dynamics,

tests/test_documentation_contract.py

@@ -300,12 +300,12 @@
     ),
     (
         "README.md",
-        "MLX 0.31's indexed-add primitive does not implement JVP",
+        "the scatter fallback requires eager, fixed indices",
         "tests/test_documentation_contract.py::test_documented_general_paths_support_jvp",
     ),
     (
         "docs/guide/performance.md",
-        "`scatter_sum` supports reverse-mode differentiation but has no forward-mode fallback",
+        "`scatter_sum(..., use_custom_kernel=False, jvp_safe=True)`",
         "tests/test_documentation_contract.py::test_documented_general_paths_support_jvp",
     ),
     (
@@ -626,14 +626,21 @@ def test_documented_general_paths_support_jvp() -> None:
 
     source = mx.arange(8, dtype=mx.float32).reshape(4, 2)
     index = mx.array([0, 1, 0, 1], dtype=mx.int32)
-    with pytest.raises(RuntimeError, match="JVP not yet implemented"):
-        mx.jvp(
-            lambda value: scatter_sum(
-                value, index, 2, use_custom_kernel=False
-            ),
-            (source,),
-            (mx.ones_like(source),),
-        )
+    (scattered,), (scatter_tangent,) = mx.jvp(
+        lambda value: scatter_sum(
+            value,
+            index,
+            2,
+            use_custom_kernel=False,
+            jvp_safe=True,
+        ),
+        (source,),
+        (mx.ones_like(source),),
+    )
+    expected_scatter = scatter_sum(source, index, 2)
+    expected_tangent = scatter_sum(mx.ones_like(source), index, 2)
+    assert _maximum_error(scattered, expected_scatter) == 0.0
+    assert _maximum_error(scatter_tangent, expected_tangent) == 0.0
 
     product = o3.FullTensorProduct("1o", "1o", use_custom_kernel=True)
     right = mx.array([[-0.5, 0.1, 0.7], [0.2, 0.4, -0.1]], dtype=mx.float32)
@@ -644,8 +651,9 @@ def test_documented_general_paths_support_jvp() -> None:
     )
     assert product_value.shape == product_tangent.shape
     assert product_value.shape == (2, product.irreps_out.dim)
-    mx.eval(harmonics_tangent, product_tangent)
+    mx.eval(harmonics_tangent, scatter_tangent, product_tangent)
     assert bool(mx.all(mx.isfinite(harmonics_tangent)))
+    assert bool(mx.all(mx.isfinite(scatter_tangent)))
     assert bool(mx.all(mx.isfinite(product_tangent)))
 
 

tests/test_graph_radial.py

@@ -119,6 +119,26 @@ def test_scatter_sum_values_multidimensional_empty_compile_and_gradient() -> Non
     gradient = mx.grad(lambda x: mx.sum(e3nn.scatter_sum(x, index, 3) ** 2))(source)
     mx.eval(gradient)
     assert bool(mx.all(mx.isfinite(gradient)))
+    safe = e3nn.scatter_sum(source, index, 3, jvp_safe=True)
+    assert _max_abs(safe - expected) == 0.0
+    safe_gradient = mx.grad(
+        lambda x: mx.sum(e3nn.scatter_sum(x, index, 3, jvp_safe=True) ** 2)
+    )(source)
+    assert _max_abs(safe_gradient - gradient) == 0.0
+    (jvp_output,), (jvp_tangent,) = mx.jvp(
+        lambda x: e3nn.scatter_sum(x, index, 3, jvp_safe=True),
+        (source,),
+        (mx.ones_like(source),),
+    )
+    assert _max_abs(jvp_output - expected) == 0.0
+    assert _max_abs(
+        jvp_tangent
+        - e3nn.scatter_sum(mx.ones_like(source), index, 3, jvp_safe=True)
+    ) == 0.0
+    compiled_safe = mx.compile(
+        lambda x: e3nn.scatter_sum(x, index, 3, jvp_safe=True)
+    )
+    assert _max_abs(compiled_safe(source) - expected) == 0.0
     empty = e3nn.scatter_sum(mx.zeros((0, 4)), mx.zeros((0,), dtype=mx.int32), 2)
     assert empty.shape == (2, 4)
 
@@ -130,6 +150,21 @@ def test_scatter_sum_validation() -> None:
         e3nn.scatter_sum(mx.ones((2, 3)), mx.array([0], dtype=mx.int32), 1)
     with pytest.raises(TypeError, match="integer"):
         e3nn.scatter_sum(mx.ones((2, 3)), mx.array([0.0, 0.0]), 1)
+    with pytest.raises(ValueError, match="requires use_custom_kernel=False"):
+        e3nn.scatter_sum(
+            mx.ones((2, 3)),
+            mx.array([0, 0], dtype=mx.int32),
+            1,
+            use_custom_kernel=True,
+            jvp_safe=True,
+        )
+    with pytest.raises(ValueError, match="0 <= index < dim_size"):
+        e3nn.scatter_sum(
+            mx.ones((2, 3)),
+            mx.array([0, 2], dtype=mx.int32),
+            2,
+            jvp_safe=True,
+        )
 
 
 @pytest.mark.mlx