Add boundschecks in numba backend

Author: aseyboldt

pytensor/link/numba/dispatch/basic.py

@@ -554,7 +554,7 @@ def numba_funcify_Subtensor(op, node, **kwargs):
         subtensor_def_src, "subtensor", {**globals(), **global_env}
     )
 
-    return numba_njit(subtensor_fn)
+    return numba_njit(subtensor_fn, boundscheck=True)
 
 
 @numba_funcify.register(IncSubtensor)
@@ -570,7 +570,7 @@ def numba_funcify_IncSubtensor(op, node, **kwargs):
         incsubtensor_def_src, "incsubtensor", {**globals(), **global_env}
     )
 
-    return numba_njit(incsubtensor_fn)
+    return numba_njit(incsubtensor_fn, boundscheck=True)
 
 
 @numba_funcify.register(AdvancedIncSubtensor1)
@@ -580,15 +580,15 @@ def numba_funcify_AdvancedIncSubtensor1(op, node, **kwargs):
 
     if set_instead_of_inc:
 
-        @numba_njit
+        @numba_njit(boundscheck=True)
         def advancedincsubtensor1_inplace(x, vals, idxs):
             for idx, val in zip(idxs, vals):
                 x[idx] = val
             return x
 
     else:
 
-        @numba_njit
+        @numba_njit(boundscheck=True)
         def advancedincsubtensor1_inplace(x, vals, idxs):
             for idx, val in zip(idxs, vals):
                 x[idx] += val

tests/link/numba/test_basic.py

@@ -373,6 +373,14 @@ def test_AdvancedSubtensor1(x, indices):
     compare_numba_and_py(out_fg, [])
 
 
+def test_AdvancedSubtensor1_out_of_bounds():
+    out_at = at_subtensor.advanced_subtensor1(np.arange(3), [4])
+    assert isinstance(out_at.owner.op, at_subtensor.AdvancedSubtensor1)
+    out_fg = FunctionGraph([], [out_at])
+    with pytest.raises(IndexError):
+        compare_numba_and_py(out_fg, [])
+
+
 @pytest.mark.parametrize(
     "x, indices",
     [