Do not use Numba objmode for supported AdvancedSubtensor operations

Author: ricardoV94

pytensor/link/numba/dispatch/basic.py

@@ -479,6 +479,9 @@ def numba_funcify_FunctionGraph(
     )
 
 
+SET_OR_INC_OPS = IncSubtensor | AdvancedIncSubtensor1 | AdvancedIncSubtensor
+
+
 def create_index_func(node, objmode=False):
     """Create a Python function that assembles and uses an index on an array."""
 
@@ -501,9 +504,7 @@ def convert_indices(indices, entry):
         else:
             raise ValueError()
 
-    set_or_inc = isinstance(
-        node.op, IncSubtensor | AdvancedIncSubtensor1 | AdvancedIncSubtensor
-    )
+    set_or_inc = isinstance(node.op, SET_OR_INC_OPS)
     index_start_idx = 1 + int(set_or_inc)
 
     input_names = [unique_names(v, force_unique=True) for v in node.inputs]
@@ -545,21 +546,6 @@ def convert_indices(indices, entry):
         index_prologue = ""
         index_body = f"z = {input_names[0]}[indices]"
 
-    if objmode:
-        output_var = node.outputs[0]
-
-        if not set_or_inc:
-            # Since `z` is being "created" while in object mode, it's
-            # considered an "outgoing" variable and needs to be manually typed
-            output_sig = f"z='{output_var.dtype}[{', '.join([':'] * output_var.ndim)}]'"
-        else:
-            output_sig = ""
-
-        index_body = f"""
-    with objmode({output_sig}):
-        {index_body}
-        """
-
     subtensor_def_src = f"""
 def {fn_name}({", ".join(input_names)}):
     {index_prologue}
@@ -572,48 +558,34 @@ def {fn_name}({", ".join(input_names)}):
 
 
 @numba_funcify.register(Subtensor)
+@numba_funcify.register(IncSubtensor)
 @numba_funcify.register(AdvancedSubtensor1)
-def numba_funcify_Subtensor(op, node, **kwargs):
-    objmode = isinstance(op, AdvancedSubtensor)
-    if objmode:
-        warnings.warn(
-            ("Numba will use object mode to allow run " "AdvancedSubtensor."),
-            UserWarning,
-        )
-
-    subtensor_def_src = create_index_func(node, objmode=objmode)
-
-    global_env = {"np": np}
-    if objmode:
-        global_env["objmode"] = numba.objmode
-
+def numba_funcify_default_subtensor(op, node, **kwargs):
+    function_name = "subtensor"
+    if isinstance(op, SET_OR_INC_OPS):
+        function_name = "setsubtensor" if op.set_instead_of_inc else "incsubtensor"
+    subtensor_def_src = create_index_func(node)
     subtensor_fn = compile_function_src(
-        subtensor_def_src, "subtensor", {**globals(), **global_env}
+        subtensor_def_src, function_name, {**globals(), "np": np}
     )
-
     return numba_njit(subtensor_fn, boundscheck=True)
 
 
-@numba_funcify.register(IncSubtensor)
-def numba_funcify_IncSubtensor(op, node, **kwargs):
-    objmode = isinstance(op, AdvancedIncSubtensor)
-    if objmode:
-        warnings.warn(
-            ("Numba will use object mode to allow run " "AdvancedIncSubtensor."),
-            UserWarning,
-        )
-
-    incsubtensor_def_src = create_index_func(node, objmode=objmode)
-
-    global_env = {"np": np}
-    if objmode:
-        global_env["objmode"] = numba.objmode
+@numba_funcify.register(AdvancedSubtensor)
+def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
+    idxs = node.inputs[1:]
+    adv_idxs_dims = [
+        idx.type.ndim
+        for idx in idxs
+        if (isinstance(idx.type, TensorType) and idx.type.ndim > 0)
+    ]
 
-    incsubtensor_fn = compile_function_src(
-        incsubtensor_def_src, "incsubtensor", {**globals(), **global_env}
-    )
+    # Numba does not support indexes with more than one dimension
+    # Nor multiple vector indexes
+    if len(adv_idxs_dims) > 1 or adv_idxs_dims[0] > 1:
+        return generate_fallback_impl(op, node, **kwargs)
 
-    return numba_njit(incsubtensor_fn, boundscheck=True)
+    return numba_funcify_default_subtensor(op, node, **kwargs)
 
 
 @numba_funcify.register(AdvancedIncSubtensor1)

tests/link/numba/test_basic.py

@@ -409,21 +409,46 @@ def test_AdvancedSubtensor1_out_of_bounds():
 
 
 @pytest.mark.parametrize(
-    "x, indices",
+    "x, indices, objmode_needed",
     [
-        (pt.as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))), ([1, 2], [2, 3])),
+        (
+            pt.as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
+            (0, [1, 2, 2, 3]),
+            False,
+        ),
+        (
+            pt.as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
+            (np.array([True, False, False])),
+            False,
+        ),
+        (
+            pt.as_tensor(np.arange(3 * 3).reshape((3, 3))),
+            (np.eye(3).astype(bool)),
+            True,
+        ),
+        (pt.as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))), ([1, 2], [2, 3]), True),
         (
             pt.as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
             ([1, 2], slice(None), [3, 4]),
+            True,
         ),
     ],
 )
-def test_AdvancedSubtensor(x, indices):
+@pytest.mark.filterwarnings("error")
+def test_AdvancedSubtensor(x, indices, objmode_needed):
     """Test NumPy's advanced indexing in more than one dimension."""
     out_pt = x[indices]
     assert isinstance(out_pt.owner.op, pt_subtensor.AdvancedSubtensor)
     out_fg = FunctionGraph([], [out_pt])
-    compare_numba_and_py(out_fg, [])
+    with (
+        pytest.warns(
+            UserWarning,
+            match="Numba will use object mode to run AdvancedSubtensor's perform method",
+        )
+        if objmode_needed
+        else contextlib.nullcontext()
+    ):
+        compare_numba_and_py(out_fg, [])
 
 
 @pytest.mark.parametrize(