WIP

Author: hameerabbasi

pixi.toml

@@ -83,10 +83,14 @@ test-finch = "ci/test_Finch.sh"
 [feature.mlir.activation.env]
 SPARSE_BACKEND = "MLIR"
 
+[feature.cuda-12.target.linux-64.pypi-dependencies]
+cupy-cuda12x = ">=13"
+
 [environments]
 test = ["test", "extra"]
 doc = ["doc", "extra"]
 mlir-dev = {features = ["test", "mlir"], no-default-feature = true}
 finch-dev = {features = ["test", "finch"], no-default-feature = true}
 notebooks = ["extra", "mlir", "finch", "notebooks"]
 barebones = {features = ["barebones"], no-default-feature = true}
+cuda-12 = ["cuda-12"]

sparse/numba_backend/_common.py

@@ -11,6 +11,7 @@
 import numpy as np
 
 from ._coo import as_coo
+from ._settings import SUPPORTED_ARRAY_TYPE
 from ._sparse_array import SparseArray
 from ._utils import (
     _zero_of_dtype,
@@ -30,6 +31,13 @@ def _is_scipy_sparse_obj(x):
     return bool(hasattr(x, "__module__") and x.__module__.startswith("scipy.sparse"))
 
 
+def _coerce_to_supported_dense(x) -> SUPPORTED_ARRAY_TYPE:
+    if isinstance(x, SUPPORTED_ARRAY_TYPE):
+        return x
+
+    return np.asarray(x)
+
+
 def _check_device(func):
     @wraps(func)
     def wrapped(*args, **kwargs):

sparse/numba_backend/_compressed/compressed.py

@@ -132,6 +132,8 @@ class GCXS(SparseArray, NDArrayOperatorsMixin):
 
     __array_priority__ = 12
 
+    __array_members__ = ("data", "indices", "indptr")
+
     def __init__(
         self,
         arg,

sparse/numba_backend/_coo/common.py

@@ -55,14 +55,15 @@ def asCOO(x, name="asCOO", check=True):
 
 
 def linear_loc(coords, shape):
+    namespace = coords.__array_namespace__()
     if shape == () and len(coords) == 0:
         # `np.ravel_multi_index` is not aware of arrays, so cannot produce a
         # sensible result here (https://github.com/numpy/numpy/issues/15690).
         # Since `coords` is an array and not a sequence, we know the correct
         # dimensions.
-        return np.zeros(coords.shape[1:], dtype=np.intp)
+        return namespace.zeros(coords.shape[1:], dtype=namespace.intp)
 
-    return np.ravel_multi_index(coords, shape)
+    return namespace.ravel_multi_index(coords, shape)
 
 
 def kron(a, b):

sparse/numba_backend/_coo/core.py

@@ -195,6 +195,8 @@ class COO(SparseArray, NDArrayOperatorsMixin):  # lgtm [py/missing-equals]
 
     __array_priority__ = 12
 
+    __array_members__ = ("data", "coords")
+
     def __init__(
         self,
         coords,
@@ -207,6 +209,8 @@ def __init__(
         fill_value=None,
         idx_dtype=None,
     ):
+        from .._common import _coerce_to_supported_dense
+
         if isinstance(coords, COO):
             self._make_shallow_copy_of(coords)
             if data is not None or shape is not None:
@@ -226,8 +230,8 @@ def __init__(
                 self.enable_caching()
             return
 
-        self.data = np.asarray(data)
-        self.coords = np.asarray(coords)
+        self.data = _coerce_to_supported_dense(data)
+        self.coords = _coerce_to_supported_dense(coords)
 
         if self.coords.ndim == 1:
             if self.coords.size == 0 and shape is not None:
@@ -236,7 +240,7 @@ def __init__(
                 self.coords = self.coords[None, :]
 
         if self.data.ndim == 0:
-            self.data = np.broadcast_to(self.data, self.coords.shape[1])
+            self.data = self._component_namespace.broadcast_to(self.data, self.coords.shape[1])
 
         if self.data.ndim != 1:
             raise ValueError("`data` must be a scalar or 1-dimensional.")
@@ -251,7 +255,9 @@ def __init__(
             shape = tuple(shape)
 
         if shape and not self.coords.size:
-            self.coords = np.zeros((len(shape) if isinstance(shape, Iterable) else 1, 0), dtype=np.intp)
+            self.coords = self._component_namespace.zeros(
+                (len(shape) if isinstance(shape, Iterable) else 1, 0), dtype=np.intp
+            )
         super().__init__(shape, fill_value=fill_value)
         if idx_dtype:
             if not can_store(idx_dtype, max(shape)):
@@ -1307,10 +1313,10 @@ def _sort_indices(self):
         """
         linear = self.linear_loc()
 
-        if (np.diff(linear) >= 0).all():  # already sorted
+        if (self._component_namespace.diff(linear) >= 0).all():  # already sorted
             return
 
-        order = np.argsort(linear, kind="mergesort")
+        order = self._component_namespace.argsort(linear, kind="mergesort")
         self.coords = self.coords[:, order]
         self.data = self.data[order]
 
@@ -1336,16 +1342,16 @@ def _sum_duplicates(self):
         # Inspired by scipy/sparse/coo.py::sum_duplicates
         # See https://github.com/scipy/scipy/blob/main/LICENSE.txt
         linear = self.linear_loc()
-        unique_mask = np.diff(linear) != 0
+        unique_mask = self._component_namespace.diff(linear) != 0
 
         if unique_mask.sum() == len(unique_mask):  # already unique
             return
 
-        unique_mask = np.append(True, unique_mask)
+        unique_mask = self._component_namespace.append(True, unique_mask)
 
         coords = self.coords[:, unique_mask]
-        (unique_inds,) = np.nonzero(unique_mask)
-        data = np.add.reduceat(self.data, unique_inds, dtype=self.data.dtype)
+        (unique_inds,) = self._component_namespace.nonzero(unique_mask)
+        data = self._component_namespace.add.reduceat(self.data, unique_inds, dtype=self.data.dtype)
 
         self.data = data
         self.coords = coords

sparse/numba_backend/_settings.py

@@ -1,3 +1,4 @@
+import importlib.util
 import os
 
 import numpy as np
@@ -17,4 +18,20 @@ def __array_function__(self, *args, **kwargs):
         return False
 
 
+def _supported_array_type() -> type[np.ndarray]:
+    try:
+        import cupy as cp
+
+        return np.ndarray | cp.ndarray
+    except ImportError:
+        return np.ndarray
+
+
+def _cupy_available() -> bool:
+    return importlib.util.find_spec("cupy") is not None
+
+
 NEP18_ENABLED = _is_nep18_enabled()
+NUMPY_DEVICE = np.asarray(5).device
+SUPPORTED_ARRAY_TYPE = _supported_array_type()
+CUPY_AVAILABLE = _cupy_available()

sparse/numba_backend/_sparse_array.py

@@ -1,4 +1,5 @@
 import contextlib
+import copy
 import operator
 import warnings
 from abc import ABCMeta, abstractmethod
@@ -27,6 +28,7 @@ class SparseArray:
     """
 
     __metaclass__ = ABCMeta
+    __array_members__: tuple[str, ...] = ()
 
     def __init__(self, shape, fill_value=None):
         if not isinstance(shape, Iterable):
@@ -45,18 +47,54 @@ def __init__(self, shape, fill_value=None):
         else:
             self.fill_value = _zero_of_dtype(self.dtype)
 
+        self.device  # noqa: B018
+
     dtype = None
 
     @property
     def device(self):
         data = getattr(self, "data", None)
-        return getattr(data, "device", "cpu")
+        device = getattr(data, "device", "cpu")
+        assert all(getattr(self, m).device == device for m in self.__array_members__)
+        return device
+
+    @property
+    def _component_namespace(self):
+        data = getattr(self, "data", None)
+        namespace = getattr(data, "__array_namespace__", np)
+        if namespace is not np:
+            namespace = namespace()
+
+        assert all(getattr(self, m).__array_namespace__() == namespace for m in self.__array_members__)
+        return namespace
 
     def to_device(self, device, /, *, stream=None):
-        if device != "cpu":
-            raise ValueError("Only `device='cpu'` is supported.")
+        if stream is not None:
+            raise NotImplementedError("Only `stream=None` is supported at the moment.")
+
+        if device == self.device:
+            return self
+
+        import cupy as cp
+
+        from ._settings import NUMPY_DEVICE
+
+        self_copy = copy.copy(self)
+        if device == NUMPY_DEVICE:
+            for member_name in self.__array_members__:
+                member_array_gpu = getattr(self, member_name)
+                member_array_cpu = cp.asnumpy(member_array_gpu)
+                setattr(self_copy, member_name, member_array_cpu)
+
+            return self_copy
+
+        for member_name in self.__array_members__:
+            member_array_source = getattr(self, member_name)
+            with cp.cuda.Device(device):
+                member_array_dest = cp.asarray(member_array_source)
+                setattr(self_copy, member_name, member_array_dest)
 
-        return self
+        return self_copy
 
     @property
     @abstractmethod

sparse/numba_backend/_utils.py

@@ -431,8 +431,11 @@ def equivalent(x, y, /, loose=False):
     >>> equivalent(np.float64(0.0), np.float64(-0.0))
     np.False_
     """
-    x = np.asarray(x)
-    y = np.asarray(y)
+    from ._common import _coerce_to_supported_dense
+
+    x = _coerce_to_supported_dense(x)
+    y = _coerce_to_supported_dense(y)
+    namespace = x.__array_namespace__()
     # Can't contain NaNs
     dt = np.result_type(x.dtype, y.dtype)
     if not any(np.issubdtype(dt, t) for t in [np.floating, np.complexfloating]):
@@ -446,9 +449,9 @@ def equivalent(x, y, /, loose=False):
         return (x == y) | ((x != x) & (y != y))
 
     if x.size == 0 or y.size == 0:
-        shape = np.broadcast_shapes(x.shape, y.shape)
-        return np.empty(shape, dtype=np.bool_)
-    x, y = np.broadcast_arrays(x[..., None], y[..., None])
+        shape = namespace.broadcast_shapes(x.shape, y.shape)
+        return namespace.empty(shape, dtype=np.bool_)
+    x, y = namespace.broadcast_arrays(x[..., None], y[..., None])
     return (x.astype(dt).view(np.uint8) == y.astype(dt).view(np.uint8)).all(axis=-1)
 
 

sparse/numba_backend/tests/test_coo.py

@@ -1913,9 +1913,3 @@ def test_to_device():
     s2 = s.to_device(s.device)
 
     assert s is s2
-
-
-def test_to_invalid_device():
-    s = sparse.random((5, 5), density=0.5)
-    with pytest.raises(ValueError, match=r"Only .* is supported."):
-        s.to_device("invalid_device")