Implement basic numba sparse Ops

Author: ricardoV94

pytensor/link/numba/dispatch/basic.py

@@ -13,7 +13,6 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.type import Type
 from pytensor.link.numba.cache import compile_numba_function_src, hash_from_pickle_dump
-from pytensor.link.numba.dispatch.sparse import CSCMatrixType, CSRMatrixType
 from pytensor.link.utils import (
     fgraph_to_python,
 )
@@ -93,12 +92,14 @@ def get_numba_type(
         numba_dtype = numba.from_dtype(dtype)
         return numba_dtype
     elif isinstance(pytensor_type, SparseTensorType):
+        from pytensor.link.numba.dispatch.sparse import CSCMatrixType, CSRMatrixType
+
         dtype = pytensor_type.numpy_dtype
-        numba_dtype = numba.from_dtype(dtype)
+        # numba_dtype = numba.from_dtype(dtype)
         if pytensor_type.format == "csr":
-            return CSRMatrixType(numba_dtype)
+            return CSRMatrixType()
         if pytensor_type.format == "csc":
-            return CSCMatrixType(numba_dtype)
+            return CSCMatrixType()
 
         raise NotImplementedError()
     else:

pytensor/link/numba/dispatch/compile_ops.py

@@ -64,6 +64,7 @@ def identity(x):
 
 @register_funcify_default_op_cache_key(DeepCopyOp)
 def numba_funcify_DeepCopyOp(op, node, **kwargs):
+    # FIXME: SparseTensorType will match on this condition, but `np.copy` doesn't work with them
     if isinstance(node.inputs[0].type, TensorType):
 
         @numba_basic.numba_njit

pytensor/link/numba/dispatch/sparse.py

@@ -0,0 +1 @@
+from pytensor.link.numba.dispatch.sparse import basic, math