Fix PyTensor-to-Numba type resolution for sparse variables

Author: brandonwillard

pytensor/link/numba/dispatch/sparse.py

@@ -13,13 +13,20 @@
     typeof_impl,
     unbox,
 )
+from numba.np.numpy_support import from_dtype
+
+from pytensor.link.numba.dispatch.basic import get_numba_type
+from pytensor.sparse.type import SparseTensorType
 
 
 class CSMatrixType(types.Type):
     """A Numba `Type` modeled after the base class `scipy.sparse.compressed._cs_matrix`."""
 
     name: str
-    instance_class: type
+
+    @staticmethod
+    def instance_class(data, indices, indptr, shape):
+        raise NotImplementedError()
 
     def __init__(self, dtype):
         self.dtype = dtype
@@ -29,6 +36,10 @@ def __init__(self, dtype):
         self.shape = types.UniTuple(types.int64, 2)
         super().__init__(self.name)
 
+    @property
+    def key(self):
+        return (self.name, self.dtype)
+
 
 make_attribute_wrapper(CSMatrixType, "data", "data")
 make_attribute_wrapper(CSMatrixType, "indices", "indices")
@@ -161,3 +172,15 @@ def ndim(inst):
         return 2
 
     return ndim
+
+
+@get_numba_type.register(SparseTensorType)
+def get_numba_type_SparseType(pytensor_type, **kwargs):
+    dtype = from_dtype(np.dtype(pytensor_type.dtype))
+
+    if pytensor_type.format == "csr":
+        return CSRMatrixType(dtype)
+    if pytensor_type.format == "csc":
+        return CSCMatrixType(dtype)
+
+    raise NotImplementedError()

tests/link/numba/test_sparse.py

@@ -1,9 +1,16 @@
 import numba
 import numpy as np
+import pytest
 import scipy as sp
 
-# Load Numba customizations
+# Make sure the Numba customizations are loaded
 import pytensor.link.numba.dispatch.sparse  # noqa: F401
+from pytensor import config
+from pytensor.sparse import Dot, SparseTensorType
+from tests.link.numba.test_basic import compare_numba_and_py
+
+
+pytestmark = pytest.mark.filterwarnings("error")
 
 
 def test_sparse_unboxing():
@@ -62,3 +69,17 @@ def test_fn(x):
     res = test_fn(x_val)
 
     assert res == 2
+
+
+def test_sparse_objmode():
+
+    x = SparseTensorType("csc", dtype=config.floatX)()
+    y = SparseTensorType("csc", dtype=config.floatX)()
+
+    out = Dot()(x, y)
+
+    x_val = sp.sparse.random(2, 2, density=0.25, dtype=config.floatX)
+    y_val = sp.sparse.random(2, 2, density=0.25, dtype=config.floatX)
+
+    with pytest.warns(UserWarning):
+        compare_numba_and_py(((x, y), (out,)), [x_val, y_val])