Repository: jax-ml/jax
Pull request #27613 (2025-03-31 14:54:40)

Add OOB checks to jax.numpy array indexing
Add OOB checks to jax.numpy array indexing



Patch:
diff --git a/jax/_src/numpy/error.py b/jax/_src/numpy/error.py
index 20dab289d..e2c23b43b 100644
--- a/jax/_src/numpy/error.py
+++ b/jax/_src/numpy/error.py
@@ -13,10 +13,11 @@
 # limitations under the License.
 
 import contextlib
-from typing import Literal
+from typing import Literal, Sequence
 
 import jax
 from jax._src import config
+from jax._src.typing import ArrayLike
 
 Category = Literal["nan", "divide", "oob"]
 
@@ -102,6 +103,56 @@ def _set_error_if_divide_by_zero(pred: jax.Array, /):
   error_check_lib.set_error_if(pred == zero, "Division by zero encountered")
 
 
+def _check_precondition_oob_gather(
+    shape: tuple[int, ...], gather_indices: ArrayLike
+) -> None:
+  """Check for out of bounds errors before calling `lax.gather`."""
+  if config.error_checking_behavior_oob.value == "ignore":
+    return
+
+  # TODO(mattjj): fix the circular import issue.
+  from jax._src import error_check as error_check_lib
+  import jax.numpy as jnp
+
+  shape = jnp.array(shape, dtype=jnp.int32)
+  error_check_lib.set_error_if(
+      jnp.logical_or(
+          jnp.min(gather_indices) < -shape,
+          jnp.max(gather_indices) >= shape,
+      ),
+      "Out of bounds encountered before calling `lax.gather`",
+  )
+
+
+def _check_precondition_oob_dynamic_slice(
+    shape: tuple[int, ...],
+    start_indices: Sequence[ArrayLike],
+    slice_sizes: list[int],
+    allow_negative_indices: list[bool],
+) -> None:
+  """Check for out of bounds errors before calling `lax.dynamic_slice`."""
+  if config.error_checking_behavior_oob.value == "ignore":
+    return
+
+  # TODO(mattjj): fix the circular import issue.
+  from jax._src import error_check as error_check_lib
+  import jax.numpy as jnp
+
+  shape = jnp.array(shape, dtype=jnp.int32)
+  start_indices = jnp.array(start_indices, dtype=jnp.int32)
+  slice_sizes = jnp.array(slice_sizes, dtype=jnp.int32)
+  allow_negative_indices = jnp.array(allow_negative_indices, dtype=jnp.bool_)
+
+  lower_bound = jnp.where(allow_negative_indices, -shape, 0)
+  error_check_lib.set_error_if(
+      jnp.logical_or(
+          jnp.minimum(start_indices, start_indices + slice_sizes) < lower_bound,
+          jnp.maximum(start_indices, start_indices + slice_sizes) >= shape,
+      ),
+      "Out of bounds encountered before calling `lax.dynamic_slice`",
+  )
+
+
 Behavior = Literal["ignore", "raise"]
 
 
diff --git a/jax/_src/numpy/indexing.py b/jax/_src/numpy/indexing.py
index 5d59bb53b..863f0c775 100644
--- a/jax/_src/numpy/indexing.py
+++ b/jax/_src/numpy/indexing.py
@@ -20,8 +20,6 @@ import operator
 import string
 from typing import Any, NamedTuple, Sequence
 
-import numpy as np
-
 import jax
 from jax import lax
 from jax._src import array
@@ -30,17 +28,19 @@ from jax._src import core
 from jax._src import dispatch
 from jax._src import dtypes
 from jax._src import errors
+from jax._src import mesh as mesh_lib
 from jax._src.api import jit
 from jax._src.lax import lax as lax_internal
 from jax._src.numpy import einsum
-from jax._src import mesh as mesh_lib
-from jax._src.pjit import auto_axes
+from jax._src.numpy import error as jnp_error
 from jax._src.numpy import lax_numpy
 from jax._src.numpy import ufuncs
 from jax._src.numpy import util
+from jax._src.pjit import auto_axes
 from jax._src.tree_util import tree_flatten
 from jax._src.typing import Array, ArrayLike, StaticScalar
-from jax._src.util import canonicalize_axis, set_module, tuple_replace, safe_zip
+from jax._src.util import canonicalize_axis, safe_zip, set_module, tuple_replace
+import numpy as np
 
 export = set_module('jax.numpy')
 
@@ -570,7 +570,7 @@ def _attempt_rewriting_take_via_slice(arr: Array, idx: Any, mode: str | None) ->
 
   idx += (arr.ndim - len(idx)) * (slice(None),)
   start_indices: Sequence[ArrayLike] = []
-  slice_sizes: Sequence[int] = []
+  slice_sizes: list[int] = []
   allow_negative_indices: list[bool] = []
 
   for ind, size in safe_zip(idx, arr.shape):
@@ -587,6 +587,7 @@ def _attempt_rewriting_take_via_slice(arr: Array, idx: Any, mode: str | None) ->
       slice_sizes.append(1)
       allow_negative_indices.append(
           not isinstance(ind, (int, np.integer)) or bool(ind < 0))
+
   # Try to use static slicing when possible.
   if all(isinstance(i, (int, np.integer)) and i >= 0 for i in start_indices):
     int_start_indices = [int(i) for i in start_indices]  # type: ignore
@@ -598,6 +599,9 @@ def _attempt_rewriting_take_via_slice(arr: Array, idx: Any, mode: str | None) ->
     # start indices to have matching types.
     if len(start_indices) > 1:
       start_indices = util.promote_dtypes(*start_indices)
+    jnp_error._check_precondition_oob_dynamic_slice(
+        arr.shape, start_indices, slice_sizes, allow_negative_indices
+    )
     arr = lax.dynamic_slice(
         arr, start_indices=start_indices, slice_sizes=slice_sizes,
         allow_negative_indices=allow_negative_indices)
@@ -640,6 +644,7 @@ def _gather(arr, treedef, static_idx, dynamic_idx, indices_are_sorted,
             unique_indices, mode, fill_value, out_sharding):
   idx = merge_static_and_dynamic_indices(treedef, static_idx, dynamic_idx)
   indexer = index_to_gather(np.shape(arr), idx)  # shared with _scatter_update
+  jnp_error._check_precondition_oob_gather(arr.shape, indexer.gather_indices)
   y = arr
 
   if fill_value is not None:


Test patch:
diff --git a/tests/jax_numpy_error_test.py b/tests/jax_numpy_error_test.py
index f2262d8b5..a38e7d550 100644
--- a/tests/jax_numpy_error_test.py
+++ b/tests/jax_numpy_error_test.py
@@ -231,6 +231,42 @@ class JaxNumpyErrorTests(jtu.JaxTestCase):
       with self.assertRaisesRegex(JaxValueError, "Division by zero"):
         error_check.raise_if_error()
 
+  @parameterized.product(jit=[True, False])
+  def test_can_raise_oob_error_take(self, jit):
+    def f(x, a):
+      return x[a]
+
+    if jit:
+      f = jax.jit(f)
+
+    x = jnp.arange(10)
+    a = jnp.int32(10)
+
+    with jnp_error.error_checking_behavior(oob="ignore"):
+      f(x, a)
+      error_check.raise_if_error()  # should not raise error
+
+    with jnp_error.error_checking_behavior(oob="raise"):
+      f(x, a)
+      with self.assertRaisesRegex(JaxValueError, "Out of bounds"):
+        error_check.raise_if_error()
+
+  def test_can_raise_oob_error_dynamic_slice(self):
+    def f(x, a):
+      return x[:, a:a+4]  # dynamic indices are non-jittable
+
+    x = jnp.arange(10).reshape(2, 5)
+    a = jnp.array(3, dtype=jnp.int32)
+
+    with jnp_error.error_checking_behavior(oob="ignore"):
+      f(x, a)
+      error_check.raise_if_error()  # should not raise error
+
+    with jnp_error.error_checking_behavior(oob="raise"):
+      f(x, a)
+      with self.assertRaisesRegex(JaxValueError, "Out of bounds"):
+        error_check.raise_if_error()
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())