Make get_numba_type dispatch on Type

Author: brandonwillard

pytensor/link/numba/dispatch/basic.py

@@ -31,6 +31,7 @@
 )
 from pytensor.scalar.basic import ScalarType
 from pytensor.scalar.math import Softplus
+from pytensor.sparse.type import SparseTensorType
 from pytensor.tensor.blas import BatchedDot
 from pytensor.tensor.math import Dot
 from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
@@ -73,14 +74,33 @@ def numba_vectorize(*args, **kwargs):
 )
 
 
-def get_numba_type(
-    pytensor_type: Type,
+@singledispatch
+def get_numba_type(pytensor_type: Type, **kwargs) -> numba.types.Type:
+    r"""Create a Numba type object for a :class:`Type`."""
+    return numba.types.pyobject
+
+
+@get_numba_type.register(SparseTensorType)
+def get_numba_type_SparseType(pytensor_type, **kwargs):
+    # This is needed to differentiate `SparseTensorType` from `TensorType`
+    return numba.types.pyobject
+
+
+@get_numba_type.register(ScalarType)
+def get_numba_type_ScalarType(pytensor_type, **kwargs):
+    dtype = np.dtype(pytensor_type.dtype)
+    numba_dtype = numba.from_dtype(dtype)
+    return numba_dtype
+
+
+@get_numba_type.register(TensorType)
+def get_numba_type_TensorType(
+    pytensor_type,
     layout: str = "A",
     force_scalar: bool = False,
     reduce_to_scalar: bool = False,
-) -> numba.types.Type:
-    r"""Create a Numba type object for a :class:`Type`.
-
+):
+    r"""
     Parameters
     ----------
     pytensor_type
@@ -92,44 +112,27 @@ def get_numba_type(
     reduce_to_scalar
         Return Numba scalars for zero dimensional :class:`TensorType`\s.
     """
-
-    if isinstance(pytensor_type, TensorType):
-        dtype = pytensor_type.numpy_dtype
-        numba_dtype = numba.from_dtype(dtype)
-        if force_scalar or (
-            reduce_to_scalar and getattr(pytensor_type, "ndim", None) == 0
-        ):
-            return numba_dtype
-        return numba.types.Array(numba_dtype, pytensor_type.ndim, layout)
-    elif isinstance(pytensor_type, ScalarType):
-        dtype = np.dtype(pytensor_type.dtype)
-        numba_dtype = numba.from_dtype(dtype)
+    dtype = pytensor_type.numpy_dtype
+    numba_dtype = numba.from_dtype(dtype)
+    if force_scalar or (reduce_to_scalar and getattr(pytensor_type, "ndim", None) == 0):
         return numba_dtype
-    else:
-        raise NotImplementedError(f"Numba type not implemented for {pytensor_type}")
+    return numba.types.Array(numba_dtype, pytensor_type.ndim, layout)
 
 
 def create_numba_signature(
-    node_or_fgraph: Union[FunctionGraph, Apply],
-    force_scalar: bool = False,
-    reduce_to_scalar: bool = False,
+    node_or_fgraph: Union[FunctionGraph, Apply], **kwargs
 ) -> numba.types.Type:
     """Create a Numba type for the signature of an `Apply` node or `FunctionGraph`."""
     input_types = []
     for inp in node_or_fgraph.inputs:
-        input_types.append(
-            get_numba_type(
-                inp.type, force_scalar=force_scalar, reduce_to_scalar=reduce_to_scalar
-            )
-        )
+        input_types.append(get_numba_type(inp.type, **kwargs))
 
     output_types = []
     for out in node_or_fgraph.outputs:
-        output_types.append(
-            get_numba_type(
-                out.type, force_scalar=force_scalar, reduce_to_scalar=reduce_to_scalar
-            )
-        )
+        output_types.append(get_numba_type(out.type, **kwargs))
+
+    if isinstance(node_or_fgraph, FunctionGraph):
+        return numba.types.Tuple(output_types)(*input_types)
 
     if len(output_types) > 1:
         return numba.types.Tuple(output_types)(*input_types)

tests/link/numba/test_basic.py

@@ -27,6 +27,7 @@
 from pytensor.link.numba.dispatch import numba_const_convert
 from pytensor.link.numba.linker import NumbaLinker
 from pytensor.raise_op import assert_op
+from pytensor.sparse.type import SparseTensorType
 from pytensor.tensor import blas
 from pytensor.tensor import subtensor as at_subtensor
 from pytensor.tensor.elemwise import Elemwise
@@ -252,26 +253,21 @@ def assert_fn(x, y):
 
 
 @pytest.mark.parametrize(
-    "v, expected, force_scalar, not_implemented",
+    "v, expected, force_scalar",
     [
-        (MyType(), None, False, True),
-        (aes.float32, numba.types.float32, False, False),
-        (at.fscalar, numba.types.Array(numba.types.float32, 0, "A"), False, False),
-        (at.fscalar, numba.types.float32, True, False),
-        (at.lvector, numba.types.int64[:], False, False),
-        (at.dmatrix, numba.types.float64[:, :], False, False),
-        (at.dmatrix, numba.types.float64, True, False),
+        (MyType(), numba.types.pyobject, False),
+        (SparseTensorType("csc", dtype=np.float64), numba.types.pyobject, False),
+        (aes.float32, numba.types.float32, False),
+        (at.fscalar, numba.types.Array(numba.types.float32, 0, "A"), False),
+        (at.fscalar, numba.types.float32, True),
+        (at.lvector, numba.types.int64[:], False),
+        (at.dmatrix, numba.types.float64[:, :], False),
+        (at.dmatrix, numba.types.float64, True),
     ],
 )
-def test_get_numba_type(v, expected, force_scalar, not_implemented):
-    cm = (
-        contextlib.suppress()
-        if not not_implemented
-        else pytest.raises(NotImplementedError)
-    )
-    with cm:
-        res = numba_basic.get_numba_type(v, force_scalar=force_scalar)
-        assert res == expected
+def test_get_numba_type(v, expected, force_scalar):
+    res = numba_basic.get_numba_type(v, force_scalar=force_scalar)
+    assert res == expected
 
 
 @pytest.mark.parametrize(