Separate interface and dispatch of numba_funcify

Author: brandonwillard

doc/extending/creating_a_numba_jax_op.rst

@@ -83,15 +83,15 @@ Here's an example for :class:`IfElse`:
        return res if n_outs > 1 else res[0]
 
 
-Step 3: Register the function with the `jax_funcify` dispatcher
+Step 3: Register the function with the `_jax_funcify` dispatcher
 ---------------------------------------------------------------
 
 With the PyTensor `Op` replicated in JAX, we’ll need to register the
 function with the PyTensor JAX `Linker`. This is done through the use of
 `singledispatch`. If you don't know how `singledispatch` works, see the
 `Python documentation <https://docs.python.org/3/library/functools.html#functools.singledispatch>`_.
 
-The relevant dispatch functions created by `singledispatch` are :func:`pytensor.link.numba.dispatch.numba_funcify` and
+The relevant dispatch functions created by `singledispatch` are :func:`pytensor.link.numba.dispatch.basic._numba_funcify` and
 :func:`pytensor.link.jax.dispatch.jax_funcify`.
 
 Here’s an example for the `Eye`\ `Op`:

pytensor/link/numba/dispatch/__init__.py

@@ -1,5 +1,9 @@
 # isort: off
-from pytensor.link.numba.dispatch.basic import numba_funcify, numba_const_convert
+from pytensor.link.numba.dispatch.basic import (
+    numba_funcify,
+    numba_const_convert,
+    numba_njit,
+)
 
 # Load dispatch specializations
 import pytensor.link.numba.dispatch.scalar

pytensor/link/numba/dispatch/basic.py

@@ -3,7 +3,7 @@
 from contextlib import contextmanager
 from functools import singledispatch
 from textwrap import dedent
-from typing import Union
+from typing import TYPE_CHECKING, Callable, Optional, Union, cast
 
 import numba
 import numba.np.unsafe.ndarray as numba_ndarray
@@ -22,6 +22,7 @@
 from pytensor.compile.ops import DeepCopyOp
 from pytensor.graph.basic import Apply, NoParams
 from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.op import Op
 from pytensor.graph.type import Type
 from pytensor.ifelse import IfElse
 from pytensor.link.utils import (
@@ -48,6 +49,10 @@
 from pytensor.tensor.type_other import MakeSlice, NoneConst
 
 
+if TYPE_CHECKING:
+    from pytensor.graph.op import StorageMapType
+
+
 def numba_njit(*args, **kwargs):
 
     if len(args) > 0 and callable(args[0]):
@@ -339,9 +344,42 @@ def numba_const_convert(data, dtype=None, **kwargs):
     return data
 
 
+def numba_funcify(obj, node=None, storage_map=None, **kwargs) -> Callable:
+    """Convert `obj` to a Numba-JITable object."""
+    return _numba_funcify(obj, node=node, storage_map=storage_map, **kwargs)
+
+
 @singledispatch
-def numba_funcify(op, node=None, storage_map=None, **kwargs):
-    """Create a Numba compatible function from an PyTensor `Op`."""
+def _numba_funcify(
+    obj,
+    node: Optional[Apply] = None,
+    storage_map: Optional["StorageMapType"] = None,
+    **kwargs,
+) -> Callable:
+    r"""Dispatch on PyTensor object types to perform Numba conversions.
+
+    Arguments
+    ---------
+    obj
+        The object used to determine the appropriate conversion function based
+        on its type.  This is generally an `Op` instance, but `FunctionGraph`\s
+        are also supported.
+    node
+        When `obj` is an `Op`, this value should be the corresponding `Apply` node.
+    storage_map
+        A storage map with, for example, the constant and `SharedVariable` values
+        of the graph being converted.
+
+    Returns
+    -------
+    A `Callable` that can be JIT-compiled in Numba using `numba.jit`.
+
+    """
+
+
+@_numba_funcify.register(Op)
+def numba_funcify_perform(op, node, storage_map=None, **kwargs) -> Callable:
+    """Create a Numba compatible function from an PyTensor `Op.perform`."""
 
     warnings.warn(
         f"Numba will use object mode to run {op}'s perform method",
@@ -392,10 +430,10 @@ def perform(*inputs):
             ret = py_perform_return(inputs)
         return ret
 
-    return perform
+    return cast(Callable, perform)
 
 
-@numba_funcify.register(OpFromGraph)
+@_numba_funcify.register(OpFromGraph)
 def numba_funcify_OpFromGraph(op, node=None, **kwargs):
 
     _ = kwargs.pop("storage_map", None)
@@ -417,7 +455,7 @@ def opfromgraph(*inputs):
     return opfromgraph
 
 
-@numba_funcify.register(FunctionGraph)
+@_numba_funcify.register(FunctionGraph)
 def numba_funcify_FunctionGraph(
     fgraph,
     node=None,
@@ -525,9 +563,9 @@ def {fn_name}({", ".join(input_names)}):
     return subtensor_def_src
 
 
-@numba_funcify.register(Subtensor)
-@numba_funcify.register(AdvancedSubtensor)
-@numba_funcify.register(AdvancedSubtensor1)
+@_numba_funcify.register(Subtensor)
+@_numba_funcify.register(AdvancedSubtensor)
+@_numba_funcify.register(AdvancedSubtensor1)
 def numba_funcify_Subtensor(op, node, **kwargs):
 
     subtensor_def_src = create_index_func(
@@ -543,8 +581,8 @@ def numba_funcify_Subtensor(op, node, **kwargs):
     return numba_njit(subtensor_fn)
 
 
-@numba_funcify.register(IncSubtensor)
-@numba_funcify.register(AdvancedIncSubtensor)
+@_numba_funcify.register(IncSubtensor)
+@_numba_funcify.register(AdvancedIncSubtensor)
 def numba_funcify_IncSubtensor(op, node, **kwargs):
 
     incsubtensor_def_src = create_index_func(
@@ -560,7 +598,7 @@ def numba_funcify_IncSubtensor(op, node, **kwargs):
     return numba_njit(incsubtensor_fn)
 
 
-@numba_funcify.register(AdvancedIncSubtensor1)
+@_numba_funcify.register(AdvancedIncSubtensor1)
 def numba_funcify_AdvancedIncSubtensor1(op, node, **kwargs):
     inplace = op.inplace
     set_instead_of_inc = op.set_instead_of_inc
@@ -593,7 +631,7 @@ def advancedincsubtensor1(x, vals, idxs):
         return advancedincsubtensor1
 
 
-@numba_funcify.register(DeepCopyOp)
+@_numba_funcify.register(DeepCopyOp)
 def numba_funcify_DeepCopyOp(op, node, **kwargs):
 
     # Scalars are apparently returned as actual Python scalar types and not
@@ -615,26 +653,26 @@ def deepcopyop(x):
     return deepcopyop
 
 
-@numba_funcify.register(MakeSlice)
-def numba_funcify_MakeSlice(op, **kwargs):
+@_numba_funcify.register(MakeSlice)
+def numba_funcify_MakeSlice(op, node, **kwargs):
     @numba_njit
     def makeslice(*x):
         return slice(*x)
 
     return makeslice
 
 
-@numba_funcify.register(Shape)
-def numba_funcify_Shape(op, **kwargs):
+@_numba_funcify.register(Shape)
+def numba_funcify_Shape(op, node, **kwargs):
     @numba_njit(inline="always")
     def shape(x):
         return np.asarray(np.shape(x))
 
     return shape
 
 
-@numba_funcify.register(Shape_i)
-def numba_funcify_Shape_i(op, **kwargs):
+@_numba_funcify.register(Shape_i)
+def numba_funcify_Shape_i(op, node, **kwargs):
     i = op.i
 
     @numba_njit(inline="always")
@@ -664,8 +702,8 @@ def codegen(context, builder, signature, args):
     return sig, codegen
 
 
-@numba_funcify.register(Reshape)
-def numba_funcify_Reshape(op, **kwargs):
+@_numba_funcify.register(Reshape)
+def numba_funcify_Reshape(op, node, **kwargs):
     ndim = op.ndim
 
     if ndim == 0:
@@ -687,7 +725,7 @@ def reshape(x, shape):
     return reshape
 
 
-@numba_funcify.register(SpecifyShape)
+@_numba_funcify.register(SpecifyShape)
 def numba_funcify_SpecifyShape(op, node, **kwargs):
     shape_inputs = node.inputs[1:]
     shape_input_names = ["shape_" + str(i) for i in range(len(shape_inputs))]
@@ -734,7 +772,7 @@ def inputs_cast(x):
     return inputs_cast
 
 
-@numba_funcify.register(Dot)
+@_numba_funcify.register(Dot)
 def numba_funcify_Dot(op, node, **kwargs):
     # Numba's `np.dot` does not support integer dtypes, so we need to cast to
     # float.
@@ -749,7 +787,7 @@ def dot(x, y):
     return dot
 
 
-@numba_funcify.register(Softplus)
+@_numba_funcify.register(Softplus)
 def numba_funcify_Softplus(op, node, **kwargs):
 
     x_dtype = np.dtype(node.inputs[0].dtype)
@@ -768,7 +806,7 @@ def softplus(x):
     return softplus
 
 
-@numba_funcify.register(Cholesky)
+@_numba_funcify.register(Cholesky)
 def numba_funcify_Cholesky(op, node, **kwargs):
     lower = op.lower
 
@@ -804,7 +842,7 @@ def cholesky(a):
     return cholesky
 
 
-@numba_funcify.register(Solve)
+@_numba_funcify.register(Solve)
 def numba_funcify_Solve(op, node, **kwargs):
 
     assume_a = op.assume_a
@@ -851,7 +889,7 @@ def solve(a, b):
     return solve
 
 
-@numba_funcify.register(BatchedDot)
+@_numba_funcify.register(BatchedDot)
 def numba_funcify_BatchedDot(op, node, **kwargs):
     dtype = node.outputs[0].type.numpy_dtype
 
@@ -872,7 +910,7 @@ def batched_dot(x, y):
 # optimizations are apparently already performed by Numba
 
 
-@numba_funcify.register(IfElse)
+@_numba_funcify.register(IfElse)
 def numba_funcify_IfElse(op, **kwargs):
     n_outs = op.n_outs
 

pytensor/link/numba/dispatch/elemwise.py

@@ -12,6 +12,7 @@
 from pytensor.graph.op import Op
 from pytensor.link.numba.dispatch import basic as numba_basic
 from pytensor.link.numba.dispatch.basic import (
+    _numba_funcify,
     create_numba_signature,
     create_tuple_creator,
     numba_funcify,
@@ -422,7 +423,7 @@ def axis_apply_fn(x):
     return axis_apply_fn
 
 
-@numba_funcify.register(Elemwise)
+@_numba_funcify.register(Elemwise)
 def numba_funcify_Elemwise(op, node, **kwargs):
 
     scalar_op_fn = numba_funcify(op.scalar_op, node=node, inline="always", **kwargs)
@@ -474,7 +475,7 @@ def {inplace_elemwise_fn_name}({input_signature_str}):
     return elemwise_fn
 
 
-@numba_funcify.register(CAReduce)
+@_numba_funcify.register(CAReduce)
 def numba_funcify_CAReduce(op, node, **kwargs):
     axes = op.axis
     if axes is None:
@@ -512,7 +513,7 @@ def numba_funcify_CAReduce(op, node, **kwargs):
     return careduce_fn
 
 
-@numba_funcify.register(DimShuffle)
+@_numba_funcify.register(DimShuffle)
 def numba_funcify_DimShuffle(op, **kwargs):
     shuffle = tuple(op.shuffle)
     transposition = tuple(op.transposition)
@@ -590,7 +591,7 @@ def dimshuffle(x):
     return dimshuffle
 
 
-@numba_funcify.register(Softmax)
+@_numba_funcify.register(Softmax)
 def numba_funcify_Softmax(op, node, **kwargs):
 
     x_at = node.inputs[0]
@@ -627,7 +628,7 @@ def softmax_py_fn(x):
     return softmax
 
 
-@numba_funcify.register(SoftmaxGrad)
+@_numba_funcify.register(SoftmaxGrad)
 def numba_funcify_SoftmaxGrad(op, node, **kwargs):
 
     sm_at = node.inputs[1]
@@ -658,7 +659,7 @@ def softmax_grad_py_fn(dy, sm):
     return softmax_grad
 
 
-@numba_funcify.register(LogSoftmax)
+@_numba_funcify.register(LogSoftmax)
 def numba_funcify_LogSoftmax(op, node, **kwargs):
 
     x_at = node.inputs[0]
@@ -692,7 +693,7 @@ def log_softmax_py_fn(x):
     return log_softmax
 
 
-@numba_funcify.register(MaxAndArgmax)
+@_numba_funcify.register(MaxAndArgmax)
 def numba_funcify_MaxAndArgmax(op, node, **kwargs):
     axis = op.axis
     x_at = node.inputs[0]