Add np.shape and ndim overloads for sparse Numba types

Author: brandonwillard

pytensor/link/numba/dispatch/sparse.py

@@ -1,3 +1,4 @@
+import numpy as np
 import scipy as sp
 import scipy.sparse
 from numba.core import cgutils, types
@@ -6,6 +7,8 @@
     box,
     make_attribute_wrapper,
     models,
+    overload,
+    overload_attribute,
     register_model,
     typeof_impl,
     unbox,
@@ -140,3 +143,21 @@ def box_matrix(typ, val, c):
     c.pyapi.decref(shape_obj)
 
     return obj
+
+
+@overload(np.shape)
+def overload_sparse_shape(x):
+    if isinstance(x, CSMatrixType):
+        return lambda x: x.shape
+
+
+@overload_attribute(CSMatrixType, "ndim")
+def overload_sparse_ndim(inst):
+
+    if not isinstance(inst, CSMatrixType):
+        return
+
+    def ndim(inst):
+        return 2
+
+    return ndim

tests/link/numba/test_sparse.py

@@ -38,3 +38,27 @@ def test_boxing(x, y):
     assert np.array_equal(res_y_val.indices, y_val.indices)
     assert np.array_equal(res_y_val.indptr, y_val.indptr)
     assert res_y_val.shape == y_val.shape
+
+
+def test_sparse_shape():
+    @numba.njit
+    def test_fn(x):
+        return np.shape(x)
+
+    x_val = sp.sparse.csr_matrix(np.eye(100))
+
+    res = test_fn(x_val)
+
+    assert res == (100, 100)
+
+
+def test_sparse_ndim():
+    @numba.njit
+    def test_fn(x):
+        return x.ndim
+
+    x_val = sp.sparse.csr_matrix(np.eye(100))
+
+    res = test_fn(x_val)
+
+    assert res == 2