Move compile ops to their own dispatch file

Author: ricardoV94

pytensor/link/numba/dispatch/__init__.py

@@ -3,6 +3,7 @@
 
 # Load dispatch specializations
 import pytensor.link.numba.dispatch.blockwise
+import pytensor.link.numba.dispatch.compile_ops
 import pytensor.link.numba.dispatch.elemwise
 import pytensor.link.numba.dispatch.extra_ops
 import pytensor.link.numba.dispatch.nlinalg

pytensor/link/numba/dispatch/basic.py

@@ -6,15 +6,10 @@
 from numba.core.errors import NumbaWarning
 from numba.cpython.unsafe.tuple import tuple_setitem  # noqa: F401
 
-from pytensor import In, config
-from pytensor.compile import NUMBA
-from pytensor.compile.builders import OpFromGraph
-from pytensor.compile.function.types import add_supervisor_to_fgraph
-from pytensor.compile.ops import DeepCopyOp, TypeCastingOp
+from pytensor import config
 from pytensor.graph.basic import Apply
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.type import Type
-from pytensor.ifelse import IfElse
 from pytensor.link.numba.dispatch.sparse import CSCMatrixType, CSRMatrixType
 from pytensor.link.utils import (
     fgraph_to_python,
@@ -280,90 +275,3 @@ def numba_funcify_FunctionGraph(
         fgraph_name=fgraph_name,
         **kwargs,
     )
-
-
-@numba_funcify.register(OpFromGraph)
-def numba_funcify_OpFromGraph(op, node=None, **kwargs):
-    _ = kwargs.pop("storage_map", None)
-
-    # Apply inner rewrites
-    # TODO: Not sure this is the right place to do this, should we have a rewrite that
-    #  explicitly triggers the optimization of the inner graphs of OpFromGraph?
-    #  The C-code defers it to the make_thunk phase
-    fgraph = op.fgraph
-    add_supervisor_to_fgraph(
-        fgraph=fgraph,
-        input_specs=[In(x, borrow=True, mutable=False) for x in fgraph.inputs],
-        accept_inplace=True,
-    )
-    NUMBA.optimizer(fgraph)
-    fgraph_fn = numba_njit(numba_funcify(op.fgraph, **kwargs))
-
-    if len(op.fgraph.outputs) == 1:
-
-        @numba_njit
-        def opfromgraph(*inputs):
-            return fgraph_fn(*inputs)[0]
-
-    else:
-
-        @numba_njit
-        def opfromgraph(*inputs):
-            return fgraph_fn(*inputs)
-
-    return opfromgraph
-
-
-@numba_funcify.register(TypeCastingOp)
-def numba_funcify_type_casting(op, **kwargs):
-    @numba_njit
-    def identity(x):
-        return x
-
-    return identity
-
-
-@numba_funcify.register(DeepCopyOp)
-def numba_funcify_DeepCopyOp(op, node, **kwargs):
-    if isinstance(node.inputs[0].type, TensorType):
-
-        @numba_njit
-        def deepcopy(x):
-            return np.copy(x)
-
-    else:
-
-        @numba_njit
-        def deepcopy(x):
-            return x
-
-    return deepcopy
-
-
-@numba_funcify.register(IfElse)
-def numba_funcify_IfElse(op, **kwargs):
-    n_outs = op.n_outs
-
-    if n_outs > 1:
-
-        @numba_njit
-        def ifelse(cond, *args):
-            if cond:
-                res = args[:n_outs]
-            else:
-                res = args[n_outs:]
-
-            return res
-
-    else:
-
-        @numba_njit
-        def ifelse(cond, *args):
-            if cond:
-                res = args[:n_outs]
-            else:
-                res = args[n_outs:]
-
-            return res[0]
-
-    return ifelse

pytensor/link/numba/dispatch/compile_ops.py

@@ -0,0 +1,117 @@
+import numpy as np
+
+from pytensor.compile.builders import OpFromGraph
+from pytensor.compile.function.types import add_supervisor_to_fgraph
+from pytensor.compile.io import In
+from pytensor.compile.mode import NUMBA
+from pytensor.compile.ops import DeepCopyOp, TypeCastingOp
+from pytensor.ifelse import IfElse
+from pytensor.link.numba.dispatch import basic as numba_basic
+from pytensor.link.numba.dispatch.basic import (
+    numba_funcify,
+    numba_njit,
+)
+from pytensor.raise_op import CheckAndRaise
+from pytensor.tensor.type import TensorType
+
+
+@numba_funcify.register(OpFromGraph)
+def numba_funcify_OpFromGraph(op, node=None, **kwargs):
+    _ = kwargs.pop("storage_map", None)
+
+    # Apply inner rewrites
+    # TODO: Not sure this is the right place to do this, should we have a rewrite that
+    #  explicitly triggers the optimization of the inner graphs of OpFromGraph?
+    #  The C-code defers it to the make_thunk phase
+    fgraph = op.fgraph
+    add_supervisor_to_fgraph(
+        fgraph=fgraph,
+        input_specs=[In(x, borrow=True, mutable=False) for x in fgraph.inputs],
+        accept_inplace=True,
+    )
+    NUMBA.optimizer(fgraph)
+    fgraph_fn = numba_njit(numba_funcify(op.fgraph, **kwargs))
+
+    if len(op.fgraph.outputs) == 1:
+
+        @numba_basic.numba_njit
+        def opfromgraph(*inputs):
+            return fgraph_fn(*inputs)[0]
+
+    else:
+
+        @numba_basic.numba_njit
+        def opfromgraph(*inputs):
+            return fgraph_fn(*inputs)
+
+    return opfromgraph
+
+
+@numba_funcify.register(TypeCastingOp)
+def numba_funcify_type_casting(op, **kwargs):
+    @numba_basic.numba_njit
+    def identity(x):
+        return x
+
+    return identity
+
+
+@numba_funcify.register(DeepCopyOp)
+def numba_funcify_DeepCopyOp(op, node, **kwargs):
+    if isinstance(node.inputs[0].type, TensorType):
+
+        @numba_basic.numba_njit
+        def deepcopy(x):
+            return np.copy(x)
+
+    else:
+
+        @numba_basic.numba_njit
+        def deepcopy(x):
+            return x
+
+    return deepcopy
+
+
+@numba_funcify.register(IfElse)
+def numba_funcify_IfElse(op, **kwargs):
+    n_outs = op.n_outs
+
+    if n_outs > 1:
+
+        @numba_basic.numba_njit
+        def ifelse(cond, *args):
+            if cond:
+                res = args[:n_outs]
+            else:
+                res = args[n_outs:]
+
+            return res
+
+    else:
+
+        @numba_basic.numba_njit
+        def ifelse(cond, *args):
+            if cond:
+                res = args[:n_outs]
+            else:
+                res = args[n_outs:]
+
+            return res[0]
+
+    return ifelse
+
+
+@numba_funcify.register(CheckAndRaise)
+def numba_funcify_CheckAndRaise(op, node, **kwargs):
+    error = op.exc_type
+    msg = op.msg
+
+    @numba_basic.numba_njit
+    def check_and_raise(x, *conditions):
+        for cond in conditions:
+            if not cond:
+                raise error(msg)
+        return x
+
+    return check_and_raise

pytensor/link/numba/dispatch/extra_ops.py

@@ -11,7 +11,6 @@
     get_numba_type,
     numba_funcify,
 )
-from pytensor.raise_op import CheckAndRaise
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.extra_ops import (
     Bartlett,
@@ -325,18 +324,3 @@ def searchsorted(a, v):
             return np.searchsorted(a, v, side)
 
     return searchsorted
-
-
-@numba_funcify.register(CheckAndRaise)
-def numba_funcify_CheckAndRaise(op, node, **kwargs):
-    error = op.exc_type
-    msg = op.msg
-
-    @numba_basic.numba_njit
-    def check_and_raise(x, *conditions):
-        for cond in conditions:
-            if not cond:
-                raise error(msg)
-        return x
-
-    return check_and_raise