Implement a copy method for Numba sparse types

Author: brandonwillard

pytensor/link/numba/dispatch/sparse.py

@@ -2,13 +2,16 @@
 import scipy as sp
 import scipy.sparse
 from numba.core import cgutils, types
+from numba.core.imputils import impl_ret_borrowed
 from numba.extending import (
     NativeValue,
     box,
+    intrinsic,
     make_attribute_wrapper,
     models,
     overload,
     overload_attribute,
+    overload_method,
     register_model,
     typeof_impl,
     unbox,
@@ -174,6 +177,42 @@ def ndim(inst):
     return ndim
 
 
+@intrinsic
+def _sparse_copy(typingctx, inst, data, indices, indptr, shape):
+    def _construct(context, builder, sig, args):
+        typ = sig.return_type
+        struct = cgutils.create_struct_proxy(typ)(context, builder)
+        _, data, indices, indptr, shape = args
+        struct.data = data
+        struct.indices = indices
+        struct.indptr = indptr
+        struct.shape = shape
+        return impl_ret_borrowed(
+            context,
+            builder,
+            sig.return_type,
+            struct._getvalue(),
+        )
+
+    sig = inst(inst, inst.data, inst.indices, inst.indptr, inst.shape)
+
+    return sig, _construct
+
+
+@overload_method(CSMatrixType, "copy")
+def overload_sparse_copy(inst):
+
+    if not isinstance(inst, CSMatrixType):
+        return
+
+    def copy(inst):
+        return _sparse_copy(
+            inst, inst.data.copy(), inst.indices.copy(), inst.indptr.copy(), inst.shape
+        )
+
+    return copy
+
+
 @get_numba_type.register(SparseTensorType)
 def get_numba_type_SparseType(pytensor_type, **kwargs):
     dtype = from_dtype(np.dtype(pytensor_type.dtype))

tests/link/numba/test_sparse.py

@@ -71,6 +71,19 @@ def test_fn(x):
     assert res == 2
 
 
+def test_sparse_copy():
+    @numba.njit
+    def test_fn(x):
+        y = x.copy()
+        return (
+            y is not x and np.all(x.data == y.data) and np.all(x.indices == y.indices)
+        )
+
+    x_val = sp.sparse.csr_matrix(np.eye(100))
+
+    assert test_fn(x_val)
+
+
 def test_sparse_objmode():
 
     x = SparseTensorType("csc", dtype=config.floatX)()