Fix JAX empty assignment indices

src/pyrecest/_backend/jax/__init__.py

@@ -8,7 +8,7 @@
 
 import jax.numpy as _jnp
 from jax import vmap
-from jax.numpy import (  # For pyrecest; For Riemannian score-based SDE
+from jax.numpy import (  # For pyrecest; For Riemannian Score-based SDE
     abs,
     all,
     allclose,
@@ -382,6 +382,10 @@ def _assignment_value_length(values):
     return 1
 
 
+def _is_empty_index_sequence(indices):
+    return _is_iterable_index(indices) and len(indices) == 0
+
+
 def _normalize_assignment_index(indices, ndim_x, axis=0):
     if _is_boolean_index(indices):
         return _jnp.asarray(indices), False, None
@@ -445,6 +449,8 @@ def assignment(x, values, indices, axis=0):
         Copy of x with the values assigned at the given indices.
     """
     x = _jnp.asarray(x)
+    if _is_empty_index_sequence(indices):
+        return x
     normalized_indices, use_vectorization, len_indices = _normalize_assignment_index(
         indices,
         x.ndim,
@@ -486,6 +492,8 @@ def assignment_by_sum(x, values, indices, axis=0):
     If a list is given, it must have the same length as indices.
     """
     x = _jnp.asarray(x)
+    if _is_empty_index_sequence(indices):
+        return x
     normalized_indices, use_vectorization, len_indices = _normalize_assignment_index(
         indices,
         x.ndim,

tests/backend_support/test_jax_assignment_empty_indices_contract.py

@@ -0,0 +1,37 @@
+"""Regression tests for raw JAX assignment helper empty-index handling."""
+
+from __future__ import annotations
+
+import importlib.util
+
+import pytest
+from tests.support.backend_runner import run_backend_code
+
+
+@pytest.mark.backend_portable
+def test_raw_jax_assignment_helpers_treat_empty_indices_as_noop():
+    if importlib.util.find_spec("jax") is None:
+        pytest.skip("JAX is not installed")
+
+    result = run_backend_code(
+        "jax",
+        """
+import importlib
+
+import numpy as np
+from pyrecest.backend import to_numpy
+
+raw_jax = importlib.import_module("pyrecest._backend.jax")
+original = raw_jax.array([1.0, 2.0, 3.0])
+
+assigned = raw_jax.assignment(original, 99.0, [])
+added = raw_jax.assignment_by_sum(original, 99.0, [])
+array_like = raw_jax.assignment([1.0, 2.0, 3.0], 99.0, [])
+
+np.testing.assert_allclose(to_numpy(assigned), [1.0, 2.0, 3.0])
+np.testing.assert_allclose(to_numpy(added), [1.0, 2.0, 3.0])
+np.testing.assert_allclose(to_numpy(array_like), [1.0, 2.0, 3.0])
+""",
+    )
+
+    assert result.returncode == 0, result.stderr