Dockerfile

Author: hameerabbasi

Dockerfile.gpu

@@ -0,0 +1,7 @@
+FROM python:3.12-slim
+
+ADD . .
+
+RUN pip install ".[cuda-12]"
+
+CMD ["python"]

pyproject.toml

@@ -7,7 +7,7 @@ name = "sparse"
 dynamic = ["version"]
 description = "Sparse n-dimensional arrays for the PyData ecosystem"
 readme = "README.md"
-dependencies = ["numpy>=1.17", "numba>=0.49"]
+dependencies = ["numpy>=1.17", "numba>=0.49", "array_api_compat>=1.11"]
 maintainers = [{ name = "Hameer Abbasi", email = "[email protected]" }]
 requires-python = ">=3.10"
 license = { file = "LICENSE" }
@@ -51,6 +51,8 @@ tests = [
     "pre-commit",
     "pytest-codspeed",
 ]
+cuda-12 = ["cupy-cuda12x"]
+cuda-11 = ["cupy-cuda11x"]
 tox = ["sparse[tests]", "tox"]
 notebooks = ["sparse[tests]", "nbmake", "matplotlib"]
 all = ["sparse[docs,tox,notebooks,mlir]", "matrepr"]

sparse/numba_backend/_common.py

@@ -262,7 +262,7 @@ def matmul(a, b):
 
     from ._settings import NUMPY_DEVICE
 
-    if a.device != NUMPY_DEVICE or b.device != NUMPY_DEVICE:
+    if getattr(a, "device", NUMPY_DEVICE) != NUMPY_DEVICE or getattr(b, "device", NUMPY_DEVICE) != NUMPY_DEVICE:
         import cupyx.scipy.sparse as cps
 
         if isinstance(a, COO):
@@ -2074,7 +2074,10 @@ def pad(array, pad_width, mode="constant", **kwargs):
     if mode.lower() != "constant":
         raise NotImplementedError(f"Mode '{mode}' is not yet supported.")
 
-    if not equivalent(kwargs.pop("constant_values", _zero_of_dtype(array.dtype)), array.fill_value):
+    if not equivalent(
+        array._component_namespace.asarray(kwargs.pop("constant_values", _zero_of_dtype(array.dtype, array.device))),
+        array.fill_value,
+    ):
         raise ValueError("constant_values can only be equal to fill value.")
 
     if kwargs:

sparse/numba_backend/_compressed/compressed.py

@@ -184,6 +184,7 @@ def __init__(
 
         self._compressed_axes = tuple(compressed_axes) if isinstance(compressed_axes, Iterable) else None
         self.fill_value = self.data.dtype.type(fill_value)
+        self.data, self.indices, self.indptr = np.asarray(self.data), np.asarray(self.indices), np.asarray(self.indptr)
 
         if prune:
             self._prune()

sparse/numba_backend/_coo/core.py

@@ -209,6 +209,8 @@ def __init__(
         fill_value=None,
         idx_dtype=None,
     ):
+        import array_api_compat
+
         from .._common import _coerce_to_supported_dense
 
         if isinstance(coords, COO):
@@ -232,6 +234,7 @@ def __init__(
 
         self.data = _coerce_to_supported_dense(data)
         self.coords = _coerce_to_supported_dense(coords)
+        xp = array_api_compat.get_namespace(self.data, self.coords)
 
         if self.coords.ndim == 1:
             if self.coords.size == 0 and shape is not None:
@@ -240,7 +243,7 @@ def __init__(
                 self.coords = self.coords[None, :]
 
         if self.data.ndim == 0:
-            self.data = self._component_namespace.broadcast_to(self.data, self.coords.shape[1])
+            self.data = xp.broadcast_to(self.data, self.coords.shape[1])
 
         if self.data.ndim != 1:
             raise ValueError("`data` must be a scalar or 1-dimensional.")
@@ -255,9 +258,7 @@ def __init__(
             shape = tuple(shape)
 
         if shape and not self.coords.size:
-            self.coords = self._component_namespace.zeros(
-                (len(shape) if isinstance(shape, Iterable) else 1, 0), dtype=np.intp
-            )
+            self.coords = xp.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)):

sparse/numba_backend/_coo/indexing.py

@@ -118,7 +118,7 @@ def getitem(x, index):
             if n != 0:
                 return x.data[mask][0]
 
-            return x.fill_value
+            return x.fill_value[()]
 
     shape = tuple(shape)
     data = x.data[mask]

sparse/numba_backend/_sparse_array.py

@@ -50,11 +50,14 @@ def __init__(self, shape, fill_value=None):
 
             self.fill_value = _zero_of_dtype(self.dtype, getattr(getattr(self, "data", None), "device", NUMPY_DEVICE))
         data = getattr(self, "data", None)
-        if data is not None:
+        if data is not None and not isinstance(data, dict):
             import array_api_compat
 
-            self.fill_value = array_api_compat.array_namespace(data).asarray(self.fill_value)
+            xp = array_api_compat.array_namespace(data)
+        else:
+            xp = np
 
+        self.fill_value = xp.asarray(self.fill_value)
         self.device  # noqa: B018
 
     dtype = None
@@ -402,7 +405,7 @@ def _gpu_ufunc(self, ufunc, method, *inputs, **kwargs):
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         from ._settings import NUMPY_DEVICE
 
-        if not all(i.device == NUMPY_DEVICE for i in inputs):
+        if not all(getattr(i, "device", NUMPY_DEVICE) == NUMPY_DEVICE for i in inputs):
             return self._gpu_ufunc(ufunc, method, *inputs, **kwargs)
         out = kwargs.pop("out", None)
         if out is not None and not all(isinstance(x, type(self)) for x in out):

sparse/numba_backend/_umath.py

@@ -514,8 +514,9 @@ def _get_fill_value(self):
         import array_api_compat
 
         from ._coo import COO
+        from ._sparse_array import SparseArray
 
-        xp = array_api_compat.array_namespace(*(a.data for a in self.args))
+        xp = array_api_compat.array_namespace(*(a.data if isinstance(a, SparseArray) else a for a in self.args))
 
         def get_zero_arg(x):
             if isinstance(x, COO):

sparse/numba_backend/_utils.py

@@ -73,7 +73,11 @@ def assert_gcxs_slicing(s, x):
 
 
 def assert_nnz(s, x):
-    fill_value = s.fill_value if hasattr(s, "fill_value") else _zero_of_dtype(s.dtype, s.device)
+    from ._settings import NUMPY_DEVICE
+
+    fill_value = (
+        s.fill_value if hasattr(s, "fill_value") else _zero_of_dtype(s.dtype, getattr(s, "device", NUMPY_DEVICE))
+    )
     assert np.sum(~equivalent(x, fill_value)) == s.nnz
 
 
@@ -442,7 +446,7 @@ def equivalent(x, y, /, loose=False):
 
     from ._common import _coerce_to_supported_dense
 
-    namespace = array_api_compat.array_namespace(x, y)
+    xp = array_api_compat.array_namespace(x, y)
     x = _coerce_to_supported_dense(x)
     y = _coerce_to_supported_dense(y)
     # Can't contain NaNs
@@ -458,9 +462,9 @@ def equivalent(x, y, /, loose=False):
         return (x == y) | ((x != x) & (y != y))
 
     if x.size == 0 or y.size == 0:
-        shape = namespace.broadcast_shapes(x.shape, y.shape)
-        return namespace.empty(shape, dtype=np.bool_)
-    x, y = namespace.broadcast_arrays(x[..., None], y[..., None])
+        shape = xp.broadcast_shapes(x.shape, y.shape)
+        return xp.empty(shape, dtype=np.bool_)
+    x, y = xp.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_compressed.py

@@ -181,8 +181,10 @@ def test_tranpose(a, b):
 @pytest.mark.parametrize("format", [sparse.COO, sparse._compressed.CSR])
 def test_to_scipy_sparse(fill_value_in, fill_value_out, format):
     s = sparse.random((3, 5), density=0.5, format=format, fill_value=fill_value_in)
-
-    if not ((fill_value_in in {0, None} and fill_value_out in {0, None}) or equivalent(fill_value_in, fill_value_out)):
+    if not (
+        (fill_value_in in {0, None} and fill_value_out in {0, None})
+        or equivalent(np.asarray(fill_value_in), np.asarray(fill_value_out))
+    ):
         with pytest.raises(ValueError, match=r"fill_value=.* but should be in .*\."):
             s.to_scipy_sparse(accept_fv=fill_value_out)
         return