From 1555a2a91dcba8faf1700e5c2cd5dc8f4b5a6b63 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Thu, 6 Mar 2025 13:48:07 +0100
Subject: [PATCH 01/21] Move cupy types to original locations

---
 typings/cupy.pyi                     | 20 --------------------
 typings/cupy/__init__.pyi            |  3 +++
 typings/cupy/_core/__init__.pyi      |  3 +++
 typings/cupy/_core/core.pyi          | 13 +++++++++++++
 typings/cupy/_creation/from_data.pyi | 14 ++++++++++++++
 5 files changed, 33 insertions(+), 20 deletions(-)
 delete mode 100644 typings/cupy.pyi
 create mode 100644 typings/cupy/__init__.pyi
 create mode 100644 typings/cupy/_core/__init__.pyi
 create mode 100644 typings/cupy/_core/core.pyi
 create mode 100644 typings/cupy/_creation/from_data.pyi

diff --git a/typings/cupy.pyi b/typings/cupy.pyi
deleted file mode 100644
index 1358682..0000000
--- a/typings/cupy.pyi
+++ /dev/null
@@ -1,20 +0,0 @@
-# SPDX-License-Identifier: MPL-2.0
-from typing import Any, Literal, Self
-
-import numpy as np
-from numpy.typing import ArrayLike, DTypeLike, NDArray
-
-class ndarray:
-    dtype: np.dtype[Any]
-    shape: tuple[int, ...]
-    def get(self) -> NDArray[Any]: ...
-    def __power__(self, other: int) -> Self: ...
-    def __array__(self) -> NDArray[Any]: ...
-
-def asarray(
-    a: ArrayLike,
-    dtype: DTypeLike | None = None,
-    order: Literal["C", "F", "A", "K", None] = None,
-    *,
-    blocking: bool = False,
-) -> ndarray: ...
diff --git a/typings/cupy/__init__.pyi b/typings/cupy/__init__.pyi
new file mode 100644
index 0000000..07276e9
--- /dev/null
+++ b/typings/cupy/__init__.pyi
@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: MPL-2.0
+from ._core import ndarray as ndarray
+from ._creation.from_data import asarray as asarray
diff --git a/typings/cupy/_core/__init__.pyi b/typings/cupy/_core/__init__.pyi
new file mode 100644
index 0000000..7f0a46a
--- /dev/null
+++ b/typings/cupy/_core/__init__.pyi
@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: MPL-2.0
+
+from .core import ndarray as ndarray
diff --git a/typings/cupy/_core/core.pyi b/typings/cupy/_core/core.pyi
new file mode 100644
index 0000000..55a1ff2
--- /dev/null
+++ b/typings/cupy/_core/core.pyi
@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: MPL-2.0
+from typing import Any, Self
+
+import numpy as np
+from numpy.typing import NDArray
+
+class ndarray:
+    dtype: np.dtype[Any]
+    shape: tuple[int, ...]
+
+    def get(self) -> NDArray[Any]: ...
+    def __power__(self, other: int) -> Self: ...
+    def __array__(self) -> NDArray[Any]: ...
diff --git a/typings/cupy/_creation/from_data.pyi b/typings/cupy/_creation/from_data.pyi
new file mode 100644
index 0000000..1f46c9c
--- /dev/null
+++ b/typings/cupy/_creation/from_data.pyi
@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: MPL-2.0
+from typing import Literal
+
+from numpy.typing import ArrayLike, DTypeLike
+
+from .._core import ndarray
+
+def asarray(
+    a: ArrayLike,
+    dtype: DTypeLike | None = None,
+    order: Literal["C", "F", "A", "K", None] = None,
+    *,
+    blocking: bool = False,
+) -> ndarray: ...

From 7fa234fd8120a254ba21adb88a7b5682c23bac33 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Sat, 8 Mar 2025 17:08:50 +0100
Subject: [PATCH 02/21] fix test utils

---
 src/fast_array_utils/conv/_to_dense.py              | 12 ++++++------
 src/fast_array_utils/stats/_mean.py                 |  2 +-
 src/fast_array_utils/stats/_mean_var.py             |  2 +-
 src/fast_array_utils/stats/_power.py                |  2 +-
 src/fast_array_utils/stats/_sum.py                  |  4 ++--
 src/fast_array_utils/types.py                       | 13 ++++++++++---
 src/testing/fast_array_utils/_array_type.py         |  6 +++++-
 tests/test_stats.py                                 |  2 +-
 tests/test_to_dense.py                              |  2 +-
 typings/cupyx/scipy/sparse/__init__.pyi             |  4 ++++
 .../cupyx/scipy/{sparse.pyi => sparse/_base.pyi}    |  0
 typings/cupyx/scipy/sparse/_compressed.pyi          |  4 ++++
 typings/cupyx/scipy/sparse/_csc.pyi                 |  7 +++++++
 typings/cupyx/scipy/sparse/_csr.pyi                 |  7 +++++++
 14 files changed, 50 insertions(+), 17 deletions(-)
 create mode 100644 typings/cupyx/scipy/sparse/__init__.pyi
 rename typings/cupyx/scipy/{sparse.pyi => sparse/_base.pyi} (100%)
 create mode 100644 typings/cupyx/scipy/sparse/_compressed.pyi
 create mode 100644 typings/cupyx/scipy/sparse/_csc.pyi
 create mode 100644 typings/cupyx/scipy/sparse/_csr.pyi

diff --git a/src/fast_array_utils/conv/_to_dense.py b/src/fast_array_utils/conv/_to_dense.py
index 2dee255..5467354 100644
--- a/src/fast_array_utils/conv/_to_dense.py
+++ b/src/fast_array_utils/conv/_to_dense.py
@@ -17,7 +17,7 @@
     MemDiskArray: TypeAlias = (
         NDArray[Any] | types.CSBase | types.H5Dataset | types.ZarrArray | types.CSDataset
     )
-    Array: TypeAlias = MemDiskArray | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+    Array: TypeAlias = MemDiskArray | types.CupyArray | types.CupyCSMatrix | types.DaskArray
 
 
 __all__ = ["to_dense"]
@@ -35,11 +35,11 @@ def to_dense(x: types.DaskArray, /, *, to_memory: Literal[True]) -> NDArray[Any]
 
 @overload
 def to_dense(
-    x: types.CupyArray | types.CupySparseMatrix, /, *, to_memory: Literal[False] = False
+    x: types.CupyArray | types.CupyCSMatrix, /, *, to_memory: Literal[False] = False
 ) -> types.CupyArray: ...
 @overload
 def to_dense(
-    x: types.CupyArray | types.CupySparseMatrix, /, *, to_memory: Literal[True]
+    x: types.CupyArray | types.CupyCSMatrix, /, *, to_memory: Literal[True]
 ) -> NDArray[Any]: ...
 
 
@@ -99,9 +99,9 @@ def _to_dense_ooc(x: types.CSDataset, /, *, to_memory: bool = False) -> NDArray[
     return to_dense(cast("types.CSBase", x.to_memory()))
 
 
-@_to_dense.register(types.CupyArray | types.CupySparseMatrix)  # type: ignore[call-overload,misc]
+@_to_dense.register(types.CupyArray | types.CupyCSMatrix)  # type: ignore[call-overload,misc]
 def _to_dense_cupy(
-    x: types.CupyArray | types.CupySparseMatrix, /, *, to_memory: bool = False
+    x: types.CupyArray | types.CupyCSMatrix, /, *, to_memory: bool = False
 ) -> NDArray[Any] | types.CupyArray:
-    x = x.toarray() if isinstance(x, types.CupySparseMatrix) else x
+    x = x.toarray() if isinstance(x, types.CupyCSMatrix) else x
     return x.get() if to_memory else x
diff --git a/src/fast_array_utils/stats/_mean.py b/src/fast_array_utils/stats/_mean.py
index a31eb6b..3f07614 100644
--- a/src/fast_array_utils/stats/_mean.py
+++ b/src/fast_array_utils/stats/_mean.py
@@ -23,7 +23,7 @@
         | types.H5Dataset
         | types.ZarrArray
         | types.CupyArray
-        | types.CupySparseMatrix
+        | types.CupyCSMatrix
     )
     Array = NonDaskArray | types.DaskArray
 
diff --git a/src/fast_array_utils/stats/_mean_var.py b/src/fast_array_utils/stats/_mean_var.py
index e1d65fa..49cf9d4 100644
--- a/src/fast_array_utils/stats/_mean_var.py
+++ b/src/fast_array_utils/stats/_mean_var.py
@@ -16,7 +16,7 @@
 
     from numpy.typing import NDArray
 
-    MemArray = NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix
+    MemArray = NDArray[Any] | types.CSBase | types.CupyArray | types.CupyCSMatrix
 
 
 __all__ = ["mean_var"]
diff --git a/src/fast_array_utils/stats/_power.py b/src/fast_array_utils/stats/_power.py
index 946074d..f1836cb 100644
--- a/src/fast_array_utils/stats/_power.py
+++ b/src/fast_array_utils/stats/_power.py
@@ -13,7 +13,7 @@
     from numpy.typing import NDArray
 
     # All supported array types except for disk ones and CSDataset
-    Array = NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+    Array = NDArray[Any] | types.CSBase | types.CupyArray | types.CupyCSMatrix | types.DaskArray
 
     _Arr = TypeVar("_Arr", bound=Array)
 
diff --git a/src/fast_array_utils/stats/_sum.py b/src/fast_array_utils/stats/_sum.py
index 76cea5a..2959934 100644
--- a/src/fast_array_utils/stats/_sum.py
+++ b/src/fast_array_utils/stats/_sum.py
@@ -23,7 +23,7 @@
         | types.H5Dataset
         | types.ZarrArray
         | types.CupyArray
-        | types.CupySparseMatrix
+        | types.CupyCSMatrix
     )
 
 
@@ -77,7 +77,7 @@ def _sum(
         assert not isinstance(x, types.CSBase | types.DaskArray)
         # np.sum supports these, but doesn’t know it. (TODO: test cupy)
         assert not isinstance(
-            x, types.ZarrArray | types.H5Dataset | types.CupyArray | types.CupySparseMatrix
+            x, types.ZarrArray | types.H5Dataset | types.CupyArray | types.CupyCSMatrix
         )
     return cast("NDArray[Any] | np.number[Any]", np.sum(x, axis=axis, dtype=dtype))
 
diff --git a/src/fast_array_utils/types.py b/src/fast_array_utils/types.py
index 23e3385..fa8c4cb 100644
--- a/src/fast_array_utils/types.py
+++ b/src/fast_array_utils/types.py
@@ -10,10 +10,14 @@
 __all__ = [
     "CSBase",
     "CupyArray",
-    "CupySparseMatrix",
+    "CupyCSCMatrix",
+    "CupyCSMatrix",
+    "CupyCSRMatrix",
     "DaskArray",
     "H5Dataset",
+    "H5Group",
     "ZarrArray",
+    "ZarrGroup",
 ]
 
 T_co = TypeVar("T_co", covariant=True)
@@ -49,9 +53,12 @@
 
 
 if TYPE_CHECKING or find_spec("cupyx"):
-    from cupyx.scipy.sparse import spmatrix as CupySparseMatrix
+    from cupyx.scipy.sparse import csc_matrix as CupyCSCMatrix
+    from cupyx.scipy.sparse import csr_matrix as CupyCSRMatrix
 else:  # pragma: no cover
-    CupySparseMatrix = type("spmatrix", (), {})
+    CupyCSCMatrix = type("csc_matrix", (), {})
+    CupyCSRMatrix = type("csr_matrix", (), {})
+CupyCSMatrix = CupyCSRMatrix | CupyCSCMatrix
 
 
 if TYPE_CHECKING:  # https://github.com/dask/dask/issues/8853
diff --git a/src/testing/fast_array_utils/_array_type.py b/src/testing/fast_array_utils/_array_type.py
index 8f6611a..22aa67e 100644
--- a/src/testing/fast_array_utils/_array_type.py
+++ b/src/testing/fast_array_utils/_array_type.py
@@ -21,7 +21,7 @@
 
     from fast_array_utils import types
 
-    InnerArrayDask = NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix
+    InnerArrayDask = NDArray[Any] | types.CSBase | types.CupyArray | types.CupyCSMatrix
     InnerArrayDisk = types.H5Dataset | types.ZarrArray
     InnerArray = InnerArrayDask | InnerArrayDisk
     Array: TypeAlias = InnerArray | types.DaskArray | types.CSDataset
@@ -219,6 +219,10 @@ def __call__(self, x: ArrayLike, /, *, dtype: DTypeLike | None = None) -> Arr:
             import cupy as cu
 
             fn = cast("ToArray[Arr]", cu.asarray)
+        elif self.cls in {types.CupyCSCMatrix, types.CupyCSRMatrix}:
+            import cupy as cu
+
+            fn = cast("ToArray[Arr]", lambda x, dtype=None: self.cls(cu.asarray(x, dtype=dtype)))  # type: ignore[call-overload,call-arg,arg-type]
         else:
             fn = cast("ToArray[Arr]", self.cls)
 
diff --git a/tests/test_stats.py b/tests/test_stats.py
index 589aeb3..749179b 100644
--- a/tests/test_stats.py
+++ b/tests/test_stats.py
@@ -113,7 +113,7 @@ def test_mean(
 )
 def test_mean_var(
     array_type: ArrayType[
-        NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+        NDArray[Any] | types.CSBase | types.CupyArray | types.CupyCSMatrix | types.DaskArray
     ],
     axis: Literal[0, 1, None],
     mean_expected: float | list[float],
diff --git a/tests/test_to_dense.py b/tests/test_to_dense.py
index 54f3ec3..05dae62 100644
--- a/tests/test_to_dense.py
+++ b/tests/test_to_dense.py
@@ -28,7 +28,7 @@ def test_to_dense(array_type: ArrayType[Array], *, to_memory: bool) -> None:
         case False, types.DaskArray():
             assert isinstance(arr, types.DaskArray)
             assert isinstance(arr._meta, np.ndarray)  # noqa: SLF001
-        case False, types.CupyArray() | types.CupySparseMatrix():
+        case False, types.CupyArray() | types.CupyCSCMatrix() | types.CupyCSRMatrix():
             assert isinstance(arr, types.CupyArray)
         case _:
             assert isinstance(arr, np.ndarray)
diff --git a/typings/cupyx/scipy/sparse/__init__.pyi b/typings/cupyx/scipy/sparse/__init__.pyi
new file mode 100644
index 0000000..052f3dc
--- /dev/null
+++ b/typings/cupyx/scipy/sparse/__init__.pyi
@@ -0,0 +1,4 @@
+# SPDX-License-Identifier: MPL-2.0
+from ._base import spmatrix as spmatrix
+from ._csc import csc_matrix as csc_matrix
+from ._csr import csr_matrix as csr_matrix
diff --git a/typings/cupyx/scipy/sparse.pyi b/typings/cupyx/scipy/sparse/_base.pyi
similarity index 100%
rename from typings/cupyx/scipy/sparse.pyi
rename to typings/cupyx/scipy/sparse/_base.pyi
diff --git a/typings/cupyx/scipy/sparse/_compressed.pyi b/typings/cupyx/scipy/sparse/_compressed.pyi
new file mode 100644
index 0000000..7147f40
--- /dev/null
+++ b/typings/cupyx/scipy/sparse/_compressed.pyi
@@ -0,0 +1,4 @@
+# SPDX-License-Identifier: MPL-2.0
+from ._base import spmatrix
+
+class _compressed_sparse_matrix(spmatrix): ...
diff --git a/typings/cupyx/scipy/sparse/_csc.pyi b/typings/cupyx/scipy/sparse/_csc.pyi
new file mode 100644
index 0000000..07fc275
--- /dev/null
+++ b/typings/cupyx/scipy/sparse/_csc.pyi
@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: MPL-2.0
+from typing import Literal
+
+from ._compressed import _compressed_sparse_matrix
+
+class csc_matrix(_compressed_sparse_matrix):
+    format: Literal["csc"] = "csc"
diff --git a/typings/cupyx/scipy/sparse/_csr.pyi b/typings/cupyx/scipy/sparse/_csr.pyi
new file mode 100644
index 0000000..26ca800
--- /dev/null
+++ b/typings/cupyx/scipy/sparse/_csr.pyi
@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: MPL-2.0
+from typing import Literal
+
+from ._compressed import _compressed_sparse_matrix
+
+class csr_matrix(_compressed_sparse_matrix):
+    format: Literal["csr"] = "csr"

From f15338ea8f0832aa4f77f735366cc285ff87dde3 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Sat, 8 Mar 2025 17:12:35 +0100
Subject: [PATCH 03/21] get

---
 tests/test_stats.py                 | 4 ++++
 typings/cupyx/scipy/sparse/_csc.pyi | 4 ++++
 typings/cupyx/scipy/sparse/_csr.pyi | 4 ++++
 3 files changed, 12 insertions(+)

diff --git a/tests/test_stats.py b/tests/test_stats.py
index 749179b..d8cee82 100644
--- a/tests/test_stats.py
+++ b/tests/test_stats.py
@@ -103,6 +103,8 @@ def test_mean(
     result = stats.mean(arr, axis=axis)  # type: ignore[arg-type]  # https://github.com/python/mypy/issues/16777
     if isinstance(result, types.DaskArray):
         result = result.compute()
+    if isinstance(result, types.CupyArray | types.CupyCSMatrix):
+        result = result.get()
     np.testing.assert_array_equal(result, expected)
 
 
@@ -156,6 +158,8 @@ def test_is_constant(
     result = stats.is_constant(x, axis=axis)
     if isinstance(result, types.DaskArray):
         result = cast("NDArray[np.bool] | bool", result.compute())
+    if isinstance(result, types.CupyArray | types.CupyCSMatrix):
+        result = result.get()
     if isinstance(expected, list):
         np.testing.assert_array_equal(expected, result)
     else:
diff --git a/typings/cupyx/scipy/sparse/_csc.pyi b/typings/cupyx/scipy/sparse/_csc.pyi
index 07fc275..dfd9e44 100644
--- a/typings/cupyx/scipy/sparse/_csc.pyi
+++ b/typings/cupyx/scipy/sparse/_csc.pyi
@@ -1,7 +1,11 @@
 # SPDX-License-Identifier: MPL-2.0
 from typing import Literal
 
+import cupy.cuda
+import scipy.sparse as sps
+
 from ._compressed import _compressed_sparse_matrix
 
 class csc_matrix(_compressed_sparse_matrix):
     format: Literal["csc"] = "csc"
+    def get(self, stream: cupy.cuda.Stream = None) -> sps.csc_matrix: ...
diff --git a/typings/cupyx/scipy/sparse/_csr.pyi b/typings/cupyx/scipy/sparse/_csr.pyi
index 26ca800..e8247a7 100644
--- a/typings/cupyx/scipy/sparse/_csr.pyi
+++ b/typings/cupyx/scipy/sparse/_csr.pyi
@@ -1,7 +1,11 @@
 # SPDX-License-Identifier: MPL-2.0
 from typing import Literal
 
+import cupy.cuda
+import scipy.sparse as sps
+
 from ._compressed import _compressed_sparse_matrix
 
 class csr_matrix(_compressed_sparse_matrix):
     format: Literal["csr"] = "csr"
+    def get(self, stream: cupy.cuda.Stream = None) -> sps.csc_matrix: ...

From afbabe53b8ba635cde15f20af6ca4914d5588d24 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Sat, 8 Mar 2025 18:19:15 +0100
Subject: [PATCH 04/21] fix to_dense

---
 src/fast_array_utils/conv/_to_dense.py |  4 +---
 tests/test_to_dense.py                 | 18 +++++++++++-------
 typings/dask/array/core.pyi            |  3 ++-
 3 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/fast_array_utils/conv/_to_dense.py b/src/fast_array_utils/conv/_to_dense.py
index 5467354..c6b438a 100644
--- a/src/fast_array_utils/conv/_to_dense.py
+++ b/src/fast_array_utils/conv/_to_dense.py
@@ -84,9 +84,7 @@ def _to_dense_cs(x: types.CSBase, /, *, to_memory: bool = False) -> NDArray[Any]
 def _to_dense_dask(
     x: types.DaskArray, /, *, to_memory: bool = False
 ) -> NDArray[Any] | types.DaskArray:
-    import dask.array as da
-
-    x = da.map_blocks(to_dense, x)
+    x = x.map_blocks(lambda x: to_dense(x, to_memory=to_memory))
     return x.compute() if to_memory else x  # type: ignore[return-value]
 
 
diff --git a/tests/test_to_dense.py b/tests/test_to_dense.py
index 05dae62..021c0e4 100644
--- a/tests/test_to_dense.py
+++ b/tests/test_to_dense.py
@@ -17,19 +17,23 @@
 
 @pytest.mark.parametrize("to_memory", [True, False], ids=["to_memory", "not_to_memory"])
 def test_to_dense(array_type: ArrayType[Array], *, to_memory: bool) -> None:
-    x = array_type([[1, 2, 3], [4, 5, 6]])
+    x = array_type([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
     if not to_memory and array_type.cls in {types.CSCDataset, types.CSRDataset}:
         with pytest.raises(ValueError, match="to_memory must be True if x is an CS{R,C}Dataset"):
             to_dense(x, to_memory=to_memory)
         return
 
     arr = to_dense(x, to_memory=to_memory)
-    match (to_memory, x):
+    assert_expected_cls(x, arr, to_memory=to_memory)
+    assert arr.shape == (2, 3)
+
+
+def assert_expected_cls(orig: Array, converted: Array, *, to_memory: bool) -> None:
+    match (to_memory, orig):
         case False, types.DaskArray():
-            assert isinstance(arr, types.DaskArray)
-            assert isinstance(arr._meta, np.ndarray)  # noqa: SLF001
+            assert isinstance(converted, types.DaskArray)
+            assert_expected_cls(orig._meta, converted._meta, to_memory=to_memory)  # noqa: SLF001
         case False, types.CupyArray() | types.CupyCSCMatrix() | types.CupyCSRMatrix():
-            assert isinstance(arr, types.CupyArray)
+            assert isinstance(converted, types.CupyArray)
         case _:
-            assert isinstance(arr, np.ndarray)
-    assert arr.shape == (2, 3)
+            assert isinstance(converted, np.ndarray)
diff --git a/typings/dask/array/core.pyi b/typings/dask/array/core.pyi
index 0372001..48f044c 100644
--- a/typings/dask/array/core.pyi
+++ b/typings/dask/array/core.pyi
@@ -20,7 +20,8 @@ _Array: TypeAlias = (
     | scipy.sparse.csr_matrix
     | scipy.sparse.csc_matrix
     | cupy.ndarray
-    | cupyx.scipy.sparse.spmatrix
+    | cupyx.scipy.sparse.csr_matrix
+    | cupyx.scipy.sparse.csc_matrix
 )
 
 class BlockView:

From 871461e60119608a51750c64d1d30b44123ba53d Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Mon, 10 Mar 2025 11:30:54 +0100
Subject: [PATCH 05/21] sum support without cupy-in-dask

---
 src/fast_array_utils/stats/_sum.py | 53 ++++++++++++++++++++----------
 tests/test_stats.py                |  5 +++
 typings/cupy/_core/core.pyi        | 10 +++++-
 3 files changed, 49 insertions(+), 19 deletions(-)

diff --git a/src/fast_array_utils/stats/_sum.py b/src/fast_array_utils/stats/_sum.py
index 2959934..45c3aed 100644
--- a/src/fast_array_utils/stats/_sum.py
+++ b/src/fast_array_utils/stats/_sum.py
@@ -16,38 +16,43 @@
     from numpy._typing._array_like import _ArrayLikeFloat_co as ArrayLike
     from numpy.typing import DTypeLike, NDArray
 
-    # all supported types except Dask and CSDataset (TODO)
-    Array = (
-        NDArray[Any]
-        | types.CSBase
-        | types.H5Dataset
-        | types.ZarrArray
-        | types.CupyArray
-        | types.CupyCSMatrix
-    )
+    # all supported types except Dask, Cupy, and CSDataset (TODO)
+    CPUArray = NDArray[Any] | types.CSBase | types.H5Dataset | types.ZarrArray
 
 
 @overload
 def sum(
-    x: ArrayLike | Array, /, *, axis: None = None, dtype: DTypeLike | None = None
+    x: ArrayLike | CPUArray | types.CupyArray | types.CupyCSMatrix,
+    /,
+    *,
+    axis: None = None,
+    dtype: DTypeLike | None = None,
 ) -> np.number[Any]: ...
 @overload
 def sum(
-    x: ArrayLike | Array, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None
+    x: ArrayLike | CPUArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None
 ) -> NDArray[Any]: ...
 @overload
+def sum(
+    x: types.CupyArray | types.CupyCSMatrix,
+    /,
+    *,
+    axis: Literal[0, 1],
+    dtype: DTypeLike | None = None,
+) -> types.CupyArray: ...
+@overload
 def sum(
     x: types.DaskArray, /, *, axis: Literal[0, 1, None] = None, dtype: DTypeLike | None = None
 ) -> types.DaskArray: ...
 
 
 def sum(
-    x: ArrayLike | Array | types.DaskArray,
+    x: ArrayLike | CPUArray | types.DaskArray,
     /,
     *,
     axis: Literal[0, 1, None] = None,
     dtype: DTypeLike | None = None,
-) -> NDArray[Any] | np.number[Any] | types.DaskArray:
+) -> NDArray[Any] | np.number[Any] | types.CupyArray | types.DaskArray:
     """Sum over both or one axis.
 
     Returns
@@ -66,22 +71,34 @@ def sum(
 
 @singledispatch
 def _sum(
-    x: ArrayLike | Array | types.DaskArray,
+    x: ArrayLike | CPUArray | types.DaskArray,
     /,
     *,
     axis: Literal[0, 1, None] = None,
     dtype: DTypeLike | None = None,
-) -> NDArray[Any] | np.number[Any] | types.DaskArray:
+) -> NDArray[Any] | np.number[Any] | types.CupyArray | types.DaskArray:
     if TYPE_CHECKING:
         # these are never passed to this fallback function, but `singledispatch` wants them
-        assert not isinstance(x, types.CSBase | types.DaskArray)
-        # np.sum supports these, but doesn’t know it. (TODO: test cupy)
         assert not isinstance(
-            x, types.ZarrArray | types.H5Dataset | types.CupyArray | types.CupyCSMatrix
+            x, types.CSBase | types.DaskArray | types.CupyArray | types.CupyCSMatrix
         )
+        # np.sum supports these, but doesn’t know it. (TODO: test cupy)
+        assert not isinstance(x, types.ZarrArray | types.H5Dataset)
     return cast("NDArray[Any] | np.number[Any]", np.sum(x, axis=axis, dtype=dtype))
 
 
+@_sum.register(types.CupyArray | types.CupyCSMatrix)  # type: ignore[call-overload,misc]
+def _sum_cupy(
+    x: types.CupyArray | types.CupyCSMatrix,
+    /,
+    *,
+    axis: Literal[0, 1, None] = None,
+    dtype: DTypeLike | None = None,
+) -> types.CupyArray | np.number[Any]:
+    arr = cast("types.CupyArray", np.sum(x, axis=axis, dtype=dtype))
+    return cast("np.number[Any]", arr.get()[()]) if axis is None else arr.squeeze()
+
+
 @_sum.register(types.CSBase)  # type: ignore[call-overload,misc]
 def _sum_cs(
     x: types.CSBase, /, *, axis: Literal[0, 1, None] = None, dtype: DTypeLike | None = None
diff --git a/tests/test_stats.py b/tests/test_stats.py
index d8cee82..559a453 100644
--- a/tests/test_stats.py
+++ b/tests/test_stats.py
@@ -62,6 +62,8 @@ def test_sum(
     axis: Literal[0, 1, None],
 ) -> None:
     np_arr = np.array([[1, 2, 3], [4, 5, 6]], dtype=dtype_in)
+    if (array_type.flags & Flags.Gpu) and np_arr.dtype.kind != "f":
+        pytest.skip("GPU arrays only support floats")
     arr = array_type(np_arr.copy())
     assert arr.dtype == dtype_in
 
@@ -73,6 +75,9 @@ def test_sum(
             sum_ = sum_.compute()  # type: ignore[assignment]
         case None, _:
             assert isinstance(sum_, np.floating | np.integer), type(sum_)
+        case 0 | 1, types.CupyArray() | types.CupyCSRMatrix() | types.CupyCSCMatrix():
+            assert isinstance(sum_, types.CupyArray), type(sum_)
+            sum_ = sum_.get()
         case 0 | 1, _:
             assert isinstance(sum_, np.ndarray), type(sum_)
         case _:
diff --git a/typings/cupy/_core/core.pyi b/typings/cupy/_core/core.pyi
index 55a1ff2..1f1d051 100644
--- a/typings/cupy/_core/core.pyi
+++ b/typings/cupy/_core/core.pyi
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: MPL-2.0
-from typing import Any, Self
+from typing import Any, Literal, Self
 
 import numpy as np
 from numpy.typing import NDArray
@@ -8,6 +8,14 @@ class ndarray:
     dtype: np.dtype[Any]
     shape: tuple[int, ...]
 
+    # cupy-specific
     def get(self) -> NDArray[Any]: ...
+
+    # operators
     def __power__(self, other: int) -> Self: ...
     def __array__(self) -> NDArray[Any]: ...
+
+    # methods
+    def squeeze(self, axis: int | None = None) -> ndarray: ...
+    def ravel(self, order: Literal["C", "F", "A", "K"] = "C") -> ndarray: ...
+    def flatten(self, order: Literal["C", "F", "A", "K"] = "C") -> ndarray: ...

From 350a21c445fb8ace83fb260605baffedb08218d0 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Mon, 10 Mar 2025 12:25:48 +0100
Subject: [PATCH 06/21] sum works

---
 src/fast_array_utils/stats/_sum.py | 8 ++++----
 tests/test_stats.py                | 2 ++
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/fast_array_utils/stats/_sum.py b/src/fast_array_utils/stats/_sum.py
index 45c3aed..8327479 100644
--- a/src/fast_array_utils/stats/_sum.py
+++ b/src/fast_array_utils/stats/_sum.py
@@ -123,13 +123,13 @@ def _sum_dask(
         raise TypeError(msg)
 
     def sum_drop_keepdims(
-        a: NDArray[Any] | types.CSBase,
+        a: NDArray[Any] | types.CSBase | types.CupyArray | types.CupyCSMatrix,
         /,
         *,
         axis: tuple[Literal[0], Literal[1]] | Literal[0, 1, None] = None,
         dtype: DTypeLike | None = None,
         keepdims: bool = False,
-    ) -> NDArray[Any]:
+    ) -> NDArray[Any] | types.CupyArray:
         del keepdims
         match axis:
             case (0 | 1 as n,):
@@ -140,8 +140,8 @@ def sum_drop_keepdims(
                 msg = f"`sum` can only sum over `axis=0|1|(0,1)` but got {axis} instead"
                 raise ValueError(msg)
         rv = sum(a, axis=axis, dtype=dtype)
-        rv = np.array(rv, ndmin=1)  # make sure rv is at least 1D
-        return rv.reshape((1, len(rv)))
+        # make sure rv is 2D
+        return np.reshape(rv, (1, 1 if rv.shape == () else len(rv)))  # type: ignore[arg-type]
 
     if dtype is None:
         # Explicitly use numpy result dtype (e.g. `NDArray[bool].sum().dtype == int64`)
diff --git a/tests/test_stats.py b/tests/test_stats.py
index 559a453..531efe9 100644
--- a/tests/test_stats.py
+++ b/tests/test_stats.py
@@ -73,6 +73,8 @@ def test_sum(
         case _, types.DaskArray():
             assert isinstance(sum_, types.DaskArray), type(sum_)
             sum_ = sum_.compute()  # type: ignore[assignment]
+            if isinstance(sum_, types.CupyArray):
+                sum_ = sum_.get()
         case None, _:
             assert isinstance(sum_, np.floating | np.integer), type(sum_)
         case 0 | 1, types.CupyArray() | types.CupyCSRMatrix() | types.CupyCSCMatrix():

From 5239a955d4b1786b60707eb9ea368ea8825d6542 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Mon, 10 Mar 2025 12:47:54 +0100
Subject: [PATCH 07/21] mean and mean_var

---
 src/fast_array_utils/stats/_mean.py        | 27 +++++++++++++---------
 src/fast_array_utils/stats/_mean_var.py    | 16 +++++++++----
 src/fast_array_utils/stats/_power.py       |  6 +++--
 tests/test_stats.py                        |  8 ++++++-
 typings/cupyx/scipy/sparse/_compressed.pyi |  9 +++++++-
 5 files changed, 46 insertions(+), 20 deletions(-)

diff --git a/src/fast_array_utils/stats/_mean.py b/src/fast_array_utils/stats/_mean.py
index 3f07614..a80881e 100644
--- a/src/fast_array_utils/stats/_mean.py
+++ b/src/fast_array_utils/stats/_mean.py
@@ -17,26 +17,31 @@
     from .. import types
 
     # all supported types except Dask and CSDataset (TODO)
-    NonDaskArray = (
-        NDArray[Any]
-        | types.CSBase
-        | types.H5Dataset
-        | types.ZarrArray
-        | types.CupyArray
-        | types.CupyCSMatrix
-    )
-    Array = NonDaskArray | types.DaskArray
+    CpuArray = NDArray[Any] | types.CSBase | types.H5Dataset | types.ZarrArray
+    Array = CpuArray | types.CupyArray | types.CupyCSMatrix | types.DaskArray
 
 
 @overload
 def mean(
-    x: NonDaskArray, /, *, axis: Literal[None] = None, dtype: DTypeLike | None = None
+    x: CpuArray | types.CupyArray | types.CupyCSMatrix,
+    /,
+    *,
+    axis: Literal[None] = None,
+    dtype: DTypeLike | None = None,
 ) -> np.number[Any]: ...
 @overload
 def mean(
-    x: NonDaskArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None
+    x: CpuArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None
 ) -> NDArray[np.number[Any]]: ...
 @overload
+def mean(
+    x: types.CupyArray | types.CupyCSMatrix,
+    /,
+    *,
+    axis: Literal[0, 1],
+    dtype: DTypeLike | None = None,
+) -> types.CupyArray: ...
+@overload
 def mean(
     x: types.DaskArray, /, *, axis: Literal[0, 1], dtype: ToDType[Any] | None = None
 ) -> types.DaskArray: ...
diff --git a/src/fast_array_utils/stats/_mean_var.py b/src/fast_array_utils/stats/_mean_var.py
index 49cf9d4..d0e41b2 100644
--- a/src/fast_array_utils/stats/_mean_var.py
+++ b/src/fast_array_utils/stats/_mean_var.py
@@ -16,7 +16,8 @@
 
     from numpy.typing import NDArray
 
-    MemArray = NDArray[Any] | types.CSBase | types.CupyArray | types.CupyCSMatrix
+    CpuArray = NDArray[Any] | types.CSBase
+    GpuArray = types.CupyArray | types.CupyCSMatrix
 
 
 __all__ = ["mean_var"]
@@ -24,12 +25,16 @@
 
 @overload
 def mean_var(
-    x: MemArray, /, *, axis: Literal[None] = None, correction: int = 0
+    x: CpuArray | GpuArray, /, *, axis: Literal[None] = None, correction: int = 0
+) -> tuple[np.float64, np.float64]: ...
+@overload
+def mean_var(
+    x: CpuArray, /, *, axis: Literal[0, 1], correction: int = 0
 ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: ...
 @overload
 def mean_var(
-    x: MemArray, /, *, axis: Literal[0, 1], correction: int = 0
-) -> tuple[np.float64, np.float64]: ...
+    x: GpuArray, /, *, axis: Literal[0, 1], correction: int = 0
+) -> tuple[types.CupyArray, types.CupyArray]: ...
 @overload
 def mean_var(
     x: types.DaskArray, /, *, axis: Literal[0, 1, None] = None, correction: int = 0
@@ -38,13 +43,14 @@ def mean_var(
 
 @no_type_check  # mypy is extremely confused
 def mean_var(
-    x: MemArray | types.DaskArray,
+    x: CpuArray | GpuArray | types.DaskArray,
     /,
     *,
     axis: Literal[0, 1, None] = None,
     correction: int = 0,
 ) -> (
     tuple[NDArray[np.float64], NDArray[np.float64]]
+    | tuple[types.CupyArray, types.CupyArray]
     | tuple[np.float64, np.float64]
     | tuple[types.DaskArray, types.DaskArray]
 ):
diff --git a/src/fast_array_utils/stats/_power.py b/src/fast_array_utils/stats/_power.py
index f1836cb..5fcd4e2 100644
--- a/src/fast_array_utils/stats/_power.py
+++ b/src/fast_array_utils/stats/_power.py
@@ -31,8 +31,10 @@ def _power(x: Array, n: int, /) -> Array:
     return x**n  # type: ignore[operator]
 
 
-@_power.register(types.CSMatrix)  # type: ignore[call-overload,misc]
-def _power_cs(x: types.CSMatrix, n: int, /) -> types.CSMatrix:
+@_power.register(types.CSMatrix | types.CupyCSMatrix)  # type: ignore[call-overload,misc]
+def _power_cs(
+    x: types.CSMatrix | types.CupyCSMatrix, n: int, /
+) -> types.CSMatrix | types.CupyCSMatrix:
     return x.power(n)
 
 
diff --git a/tests/test_stats.py b/tests/test_stats.py
index 531efe9..b72d747 100644
--- a/tests/test_stats.py
+++ b/tests/test_stats.py
@@ -104,6 +104,8 @@ def test_mean(
     array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: float | list[float]
 ) -> None:
     np_arr = np.array([[1, 2, 3], [4, 5, 6]])
+    if (array_type.flags & Flags.Gpu) and np_arr.dtype.kind != "f":
+        pytest.skip("GPU arrays only support floats")
     np.testing.assert_array_equal(np.mean(np_arr, axis=axis), expected)
 
     arr = array_type(np_arr)
@@ -128,12 +130,16 @@ def test_mean_var(
     mean_expected: float | list[float],
     var_expected: float | list[float],
 ) -> None:
-    np_arr = np.array([[1, 2, 3], [4, 5, 6]])
+    np_arr = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float64)
     np.testing.assert_array_equal(np.mean(np_arr, axis=axis), mean_expected)
     np.testing.assert_array_equal(np.var(np_arr, axis=axis, correction=1), var_expected)
 
     arr = array_type(np_arr)
     mean, var = stats.mean_var(arr, axis=axis, correction=1)
+    if isinstance(mean, types.DaskArray) and isinstance(var, types.DaskArray):
+        mean, var = mean.compute(), var.compute()  # type: ignore[assignment]
+    if isinstance(mean, types.CupyArray) and isinstance(var, types.CupyArray):
+        mean, var = mean.get(), var.get()
     np.testing.assert_array_equal(mean, mean_expected)  # type: ignore[arg-type]
     np.testing.assert_array_almost_equal(var, var_expected)  # type: ignore[arg-type]
 
diff --git a/typings/cupyx/scipy/sparse/_compressed.pyi b/typings/cupyx/scipy/sparse/_compressed.pyi
index 7147f40..7e80a83 100644
--- a/typings/cupyx/scipy/sparse/_compressed.pyi
+++ b/typings/cupyx/scipy/sparse/_compressed.pyi
@@ -1,4 +1,11 @@
 # SPDX-License-Identifier: MPL-2.0
+from typing import Literal, Self
+
+from numpy.typing import DTypeLike
+
 from ._base import spmatrix
 
-class _compressed_sparse_matrix(spmatrix): ...
+class _compressed_sparse_matrix(spmatrix):
+    format: Literal["csr", "csc"]
+
+    def power(self, n: int, dtype: DTypeLike | None = None) -> Self: ...

From 2577cd3a3dc9661b1fbf3fb7c7432823879d4e83 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Mon, 10 Mar 2025 13:14:45 +0100
Subject: [PATCH 08/21] add workflow

---
 .github/workflows/ci-gpu.yml | 68 ++++++++++++++++++++++++++++++++++++
 1 file changed, 68 insertions(+)
 create mode 100644 .github/workflows/ci-gpu.yml

diff --git a/.github/workflows/ci-gpu.yml b/.github/workflows/ci-gpu.yml
new file mode 100644
index 0000000..765cf70
--- /dev/null
+++ b/.github/workflows/ci-gpu.yml
@@ -0,0 +1,68 @@
+name: GPU-CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    types:
+      - labeled
+      - opened
+      - synchronize
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  check:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: flying-sheep/check@v1
+        with:
+          success: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'run-gpu-ci') }}
+  test:
+    name: GPU Tests
+    needs: check
+    runs-on: "cirun-aws-gpu--${{ github.run_id }}"
+    timeout-minutes: 30
+    defaults:
+      run:
+        shell: bash -el {0}
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Nvidia SMI sanity check
+        run: nvidia-smi
+
+      - name: Install Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install UV
+        uses: hynek/setup-cached-uv@v2
+        with:
+          cache-dependency-path: pyproject.toml
+
+      - name: Set UV Timeout
+        env:
+          UV_HTTP_TIMEOUT: 120
+        run: echo "UV_HTTP_TIMEOUT is set to $UV_HTTP_TIMEOUT"
+
+      - name: Install package
+        run: uv pip install --system -e .[test,full] cupy-cuda12x --extra-index-url=https://pypi.nvidia.com --index-strategy=unsafe-best-match
+
+      - name: List installed packages
+        run: uv pip list
+
+      - name: Run tests
+        run: pytest
+
+      - name: Remove 'run-gpu-ci' Label
+        if: always()
+        uses: actions-ecosystem/action-remove-labels@v1
+        with:
+          labels: 'run-gpu-ci'
+          github_token: ${{ secrets.GITHUB_TOKEN }}

From b2a663396973c84e58d77d790b5c530bd845cc14 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Mon, 10 Mar 2025 13:19:59 +0100
Subject: [PATCH 09/21] cirun yml

---
 .cirun.yml | 11 +++++++++++
 1 file changed, 11 insertions(+)
 create mode 100644 .cirun.yml

diff --git a/.cirun.yml b/.cirun.yml
new file mode 100644
index 0000000..263d0fa
--- /dev/null
+++ b/.cirun.yml
@@ -0,0 +1,11 @@
+runners:
+  - name: aws-gpu-runner
+    cloud: aws
+    instance_type: g4dn.xlarge
+    machine_image: ami-067a4ba2816407ee9
+    region: eu-north-1
+    preemptible:
+      - true
+      - false
+    labels:
+      - cirun-aws-gpu

From 1cbe91d1d006901d526ce908324014e90449928a Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Mon, 10 Mar 2025 13:51:40 +0100
Subject: [PATCH 10/21] cupy array support for is_constant

---
 src/fast_array_utils/stats/_is_constant.py  | 25 +++++++++++-------
 src/testing/fast_array_utils/_array_type.py | 10 +++++--
 tests/test_sparse.py                        |  2 +-
 tests/test_stats.py                         |  7 +++--
 typings/cupy/_core/core.pyi                 | 29 +++++++++++++++++----
 typings/cupyx/scipy/sparse/_compressed.pyi  | 15 ++++++++++-
 6 files changed, 66 insertions(+), 22 deletions(-)

diff --git a/src/fast_array_utils/stats/_is_constant.py b/src/fast_array_utils/stats/_is_constant.py
index 86ee6b4..7cc153b 100644
--- a/src/fast_array_utils/stats/_is_constant.py
+++ b/src/fast_array_utils/stats/_is_constant.py
@@ -20,16 +20,23 @@
 
 
 @overload
-def is_constant(a: types.DaskArray, /, *, axis: Literal[0, 1, None] = None) -> types.DaskArray: ...
-@overload
-def is_constant(a: NDArray[Any] | types.CSBase, /, *, axis: None = None) -> bool: ...
+def is_constant(
+    a: NDArray[Any] | types.CSBase | types.CupyArray, /, *, axis: None = None
+) -> bool: ...
 @overload
 def is_constant(a: NDArray[Any] | types.CSBase, /, *, axis: Literal[0, 1]) -> NDArray[np.bool]: ...
+@overload
+def is_constant(a: types.CupyArray, /, *, axis: Literal[0, 1]) -> types.CupyArray: ...
+@overload
+def is_constant(a: types.DaskArray, /, *, axis: Literal[0, 1, None] = None) -> types.DaskArray: ...
 
 
 def is_constant(
-    a: NDArray[Any] | types.CSBase | types.DaskArray, /, *, axis: Literal[0, 1, None] = None
-) -> bool | NDArray[np.bool] | types.DaskArray:
+    a: NDArray[Any] | types.CSBase | types.CupyArray | types.DaskArray,
+    /,
+    *,
+    axis: Literal[0, 1, None] = None,
+) -> bool | NDArray[np.bool] | types.CupyArray | types.DaskArray:
     """Check whether values in array are constant.
 
     Params
@@ -69,10 +76,10 @@ def _is_constant(
     raise NotImplementedError
 
 
-@_is_constant.register(np.ndarray)
+@_is_constant.register(np.ndarray | types.CupyArray)  # type: ignore[call-overload,misc]
 def _is_constant_ndarray(
-    a: NDArray[Any], /, *, axis: Literal[0, 1, None] = None
-) -> bool | NDArray[np.bool]:
+    a: NDArray[Any] | types.CupyArray, /, *, axis: Literal[0, 1, None] = None
+) -> bool | NDArray[np.bool] | types.CupyArray:
     # Should eventually support nd, not now.
     match axis:
         case None:
@@ -83,7 +90,7 @@ def _is_constant_ndarray(
             return _is_constant_rows(a)
 
 
-def _is_constant_rows(a: NDArray[Any]) -> NDArray[np.bool]:
+def _is_constant_rows(a: NDArray[Any] | types.CupyArray) -> NDArray[np.bool] | types.CupyArray:
     b = np.broadcast_to(a[:, 0][:, np.newaxis], a.shape)
     return cast(NDArray[np.bool], (a == b).all(axis=1))
 
diff --git a/src/testing/fast_array_utils/_array_type.py b/src/testing/fast_array_utils/_array_type.py
index 22aa67e..cfd1f36 100644
--- a/src/testing/fast_array_utils/_array_type.py
+++ b/src/testing/fast_array_utils/_array_type.py
@@ -173,9 +173,15 @@ def random(
                     ),
                 )
             case "cupy", "ndarray", None:
-                raise NotImplementedError
+                return self(gen.random(shape, dtype=dtype or np.float64))
             case "cupyx.scipy.sparse", ("csr_matrix" | "csc_matrix") as cls_name, None:
-                raise NotImplementedError
+                import cupy as cu
+
+                fmt = cast('Literal["csr", "csc"]', cls_name[:3])
+                m = random_mat(shape, density=density, format=fmt, container="matrix", dtype=dtype)
+                d, i, p = tuple(cu.asarray(p) for p in (m.data, m.indices, m.indptr))
+                cls = cast("type[types.CupyCSMatrix]", self.cls)
+                return cast("Arr", cls((d, i, p), shape=shape))
             case "dask.array", "Array", _:
                 import dask.array as da
 
diff --git a/tests/test_sparse.py b/tests/test_sparse.py
index 2ac3089..7e0950e 100644
--- a/tests/test_sparse.py
+++ b/tests/test_sparse.py
@@ -54,7 +54,7 @@ def test_to_dense(
 
 
 @pytest.mark.benchmark
-@pytest.mark.array_type(select=Flags.Sparse, skip=Flags.Dask | Flags.Disk)
+@pytest.mark.array_type(select=Flags.Sparse, skip=Flags.Dask | Flags.Disk | Flags.Gpu)
 @pytest.mark.parametrize("order", ["C", "F"])
 def test_to_dense_benchmark(
     benchmark: BenchmarkFixture,
diff --git a/tests/test_stats.py b/tests/test_stats.py
index b72d747..5e83239 100644
--- a/tests/test_stats.py
+++ b/tests/test_stats.py
@@ -144,8 +144,7 @@ def test_mean_var(
     np.testing.assert_array_almost_equal(var, var_expected)  # type: ignore[arg-type]
 
 
-# TODO(flying-sheep): enable for GPU  # noqa: TD003
-@pytest.mark.array_type(skip=Flags.Disk | Flags.Gpu)
+@pytest.mark.array_type(skip=Flags.Disk)
 @pytest.mark.parametrize(
     ("axis", "expected"),
     [
@@ -167,7 +166,7 @@ def test_is_constant(
         [0, 0, 1, 0],
         [0, 0, 0, 0],
     ]
-    x = array_type(x_data)
+    x = array_type(x_data, dtype=np.float64)
     result = stats.is_constant(x, axis=axis)
     if isinstance(result, types.DaskArray):
         result = cast("NDArray[np.bool] | bool", result.compute())
@@ -184,7 +183,7 @@ def test_dask_constant_blocks(
     dask_viz: Callable[[object], None], array_type: ArrayType[types.DaskArray, Any]
 ) -> None:
     """Tests if is_constant works if each chunk is individually constant."""
-    x_np = np.repeat(np.repeat(np.arange(4).reshape(2, 2), 2, axis=0), 2, axis=1)
+    x_np = np.repeat(np.repeat(np.arange(4, dtype=np.float64).reshape(2, 2), 2, axis=0), 2, axis=1)
     x = array_type(x_np)
     assert x.blocks.shape == (2, 2)
     assert all(stats.is_constant(block).compute() for block in x.blocks.ravel())
diff --git a/typings/cupy/_core/core.pyi b/typings/cupy/_core/core.pyi
index 1f1d051..e95e036 100644
--- a/typings/cupy/_core/core.pyi
+++ b/typings/cupy/_core/core.pyi
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: MPL-2.0
-from typing import Any, Literal, Self
+from types import EllipsisType
+from typing import Any, Literal, Self, overload
 
 import numpy as np
 from numpy.typing import NDArray
@@ -12,10 +13,28 @@ class ndarray:
     def get(self) -> NDArray[Any]: ...
 
     # operators
-    def __power__(self, other: int) -> Self: ...
     def __array__(self) -> NDArray[Any]: ...
+    def __getitem__(  # never returns scalars
+        self, index: int | slice | EllipsisType | tuple[int | slice | EllipsisType | None, ...]
+    ) -> Self: ...
+    def __eq__(self, value: object) -> ndarray: ...  # type: ignore[override]
+    def __power__(self, other: int) -> Self: ...
 
     # methods
-    def squeeze(self, axis: int | None = None) -> ndarray: ...
-    def ravel(self, order: Literal["C", "F", "A", "K"] = "C") -> ndarray: ...
-    def flatten(self, order: Literal["C", "F", "A", "K"] = "C") -> ndarray: ...
+    @property
+    def T(self) -> Self: ...  # noqa: N802
+    @overload
+    def all(self, axis: None = None) -> np.bool: ...
+    @overload
+    def all(self, axis: int) -> ndarray: ...
+    def squeeze(self, axis: int | None = None) -> Self: ...
+    def ravel(self, order: Literal["C", "F", "A", "K"] = "C") -> Self: ...
+    def flatten(self, order: Literal["C", "F", "A", "K"] = "C") -> Self: ...
+    @property
+    def flat(self) -> _FlatIter: ...
+
+class _FlatIter:
+    def __next__(self) -> np.float32 | np.float64: ...
+    def __iter__(self) -> _FlatIter: ...
+    def __len__(self) -> int: ...
+    def __getitem__(self, index: int) -> np.float32 | np.float64: ...
diff --git a/typings/cupyx/scipy/sparse/_compressed.pyi b/typings/cupyx/scipy/sparse/_compressed.pyi
index 7e80a83..a414676 100644
--- a/typings/cupyx/scipy/sparse/_compressed.pyi
+++ b/typings/cupyx/scipy/sparse/_compressed.pyi
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: MPL-2.0
-from typing import Literal, Self
+from typing import Literal, Self, overload
 
+from cupy import ndarray
 from numpy.typing import DTypeLike
 
 from ._base import spmatrix
@@ -8,4 +9,16 @@ from ._base import spmatrix
 class _compressed_sparse_matrix(spmatrix):
     format: Literal["csr", "csc"]
 
+    @overload
+    def __init__(self, arg1: ndarray | spmatrix) -> None: ...
+    @overload
+    def __init__(self, arg1: tuple[int, int], *, dtype: DTypeLike | None = None) -> None: ...
+    @overload
+    def __init__(self, arg1: tuple[ndarray, tuple[ndarray, ndarray]]) -> None: ...
+    @overload
+    def __init__(
+        self, arg1: tuple[ndarray, ndarray, ndarray], shape: tuple[int, int] | None = None
+    ) -> None: ...
+
+    # methods
     def power(self, n: int, dtype: DTypeLike | None = None) -> Self: ...

From 53798d0e4233dd9f065b5b8e45cfdc4933b886bf Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Mon, 10 Mar 2025 13:52:58 +0100
Subject: [PATCH 11/21] color tests

---
 .editorconfig                |  2 +-
 .github/workflows/ci-gpu.yml | 10 +++++-----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/.editorconfig b/.editorconfig
index 0e3c023..0fd9346 100644
--- a/.editorconfig
+++ b/.editorconfig
@@ -9,6 +9,6 @@ max_line_length = 88
 indent_size = 4
 indent_style = space
 
-[*.toml]
+[*.{toml,yml,yaml}]
 indent_size = 2
 max_line_length = 120
diff --git a/.github/workflows/ci-gpu.yml b/.github/workflows/ci-gpu.yml
index 765cf70..402eaf1 100644
--- a/.github/workflows/ci-gpu.yml
+++ b/.github/workflows/ci-gpu.yml
@@ -9,6 +9,11 @@ on:
       - opened
       - synchronize
 
+env:
+  PYTEST_ADDOPTS: "-v --color=yes"
+  FORCE_COLOR: "1"
+  UV_HTTP_TIMEOUT: 120
+
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
   cancel-in-progress: true
@@ -46,11 +51,6 @@ jobs:
         with:
           cache-dependency-path: pyproject.toml
 
-      - name: Set UV Timeout
-        env:
-          UV_HTTP_TIMEOUT: 120
-        run: echo "UV_HTTP_TIMEOUT is set to $UV_HTTP_TIMEOUT"
-
       - name: Install package
         run: uv pip install --system -e .[test,full] cupy-cuda12x --extra-index-url=https://pypi.nvidia.com --index-strategy=unsafe-best-match
 

From 4d4ed71c335f9a783039e11ca4f323829f4da849 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Mon, 10 Mar 2025 14:01:05 +0100
Subject: [PATCH 12/21] skip remaining tests

---
 tests/test_sparse.py | 2 +-
 tests/test_stats.py  | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/tests/test_sparse.py b/tests/test_sparse.py
index 7e0950e..5554315 100644
--- a/tests/test_sparse.py
+++ b/tests/test_sparse.py
@@ -39,7 +39,7 @@ def dtype(request: pytest.FixtureRequest) -> type[np.float32 | np.float64]:
     return cast("type[np.float32 | np.float64]", request.param)
 
 
-@pytest.mark.array_type(select=Flags.Sparse, skip=Flags.Dask | Flags.Disk)
+@pytest.mark.array_type(select=Flags.Sparse, skip=Flags.Dask | Flags.Disk | Flags.Gpu)
 @pytest.mark.parametrize("order", ["C", "F"])
 def test_to_dense(
     array_type: ArrayType[CSBase, None],
diff --git a/tests/test_stats.py b/tests/test_stats.py
index 5e83239..9801e84 100644
--- a/tests/test_stats.py
+++ b/tests/test_stats.py
@@ -37,6 +37,7 @@ def __call__(  # noqa: D102
 
 # can’t select these using a category filter
 ATS_SPARSE_DS = {at for at in SUPPORTED_TYPES if at.mod == "anndata.abc"}
+ATS_CUPY_SPARSE = {at for at in SUPPORTED_TYPES if "cupyx.scipy" in str(at)}
 
 
 @pytest.fixture(scope="session", params=[0, 1, None])
@@ -144,7 +145,7 @@ def test_mean_var(
     np.testing.assert_array_almost_equal(var, var_expected)  # type: ignore[arg-type]
 
 
-@pytest.mark.array_type(skip=Flags.Disk)
+@pytest.mark.array_type(skip={Flags.Disk, *ATS_CUPY_SPARSE})
 @pytest.mark.parametrize(
     ("axis", "expected"),
     [
@@ -178,7 +179,7 @@ def test_is_constant(
         assert expected is result
 
 
-@pytest.mark.array_type(Flags.Dask)
+@pytest.mark.array_type(Flags.Dask, skip=ATS_CUPY_SPARSE)
 def test_dask_constant_blocks(
     dask_viz: Callable[[object], None], array_type: ArrayType[types.DaskArray, Any]
 ) -> None:

From 2144a177ab271e0a7f26216098fea1a18c0b9b53 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Mon, 10 Mar 2025 14:02:11 +0100
Subject: [PATCH 13/21] fix type

---
 src/fast_array_utils/stats/_is_constant.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/fast_array_utils/stats/_is_constant.py b/src/fast_array_utils/stats/_is_constant.py
index 7cc153b..32e39bb 100644
--- a/src/fast_array_utils/stats/_is_constant.py
+++ b/src/fast_array_utils/stats/_is_constant.py
@@ -71,8 +71,11 @@ def is_constant(
 
 @singledispatch
 def _is_constant(
-    a: NDArray[Any] | types.CSBase | types.DaskArray, /, *, axis: Literal[0, 1, None] = None
-) -> bool | NDArray[np.bool] | types.DaskArray:  # pragma: no cover
+    a: NDArray[Any] | types.CSBase | types.CupyArray | types.DaskArray,
+    /,
+    *,
+    axis: Literal[0, 1, None] = None,
+) -> bool | NDArray[np.bool] | types.CupyArray | types.DaskArray:  # pragma: no cover
     raise NotImplementedError
 
 

From bbf256bd9787af3eb86da70ba74e1e490a68f6c3 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Mon, 10 Mar 2025 14:08:46 +0100
Subject: [PATCH 14/21] fix docs

---
 src/fast_array_utils/stats/_sum.py | 31 +++++++++++++++++++++++-------
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/fast_array_utils/stats/_sum.py b/src/fast_array_utils/stats/_sum.py
index 8327479..a5937dc 100644
--- a/src/fast_array_utils/stats/_sum.py
+++ b/src/fast_array_utils/stats/_sum.py
@@ -16,13 +16,16 @@
     from numpy._typing._array_like import _ArrayLikeFloat_co as ArrayLike
     from numpy.typing import DTypeLike, NDArray
 
-    # all supported types except Dask, Cupy, and CSDataset (TODO)
-    CPUArray = NDArray[Any] | types.CSBase | types.H5Dataset | types.ZarrArray
-
 
 @overload
 def sum(
-    x: ArrayLike | CPUArray | types.CupyArray | types.CupyCSMatrix,
+    x: ArrayLike
+    | NDArray[Any]
+    | types.CSBase
+    | types.H5Dataset
+    | types.ZarrArray
+    | types.CupyArray
+    | types.CupyCSMatrix,
     /,
     *,
     axis: None = None,
@@ -30,7 +33,11 @@ def sum(
 ) -> np.number[Any]: ...
 @overload
 def sum(
-    x: ArrayLike | CPUArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None
+    x: ArrayLike | NDArray[Any] | types.CSBase | types.H5Dataset | types.ZarrArray,
+    /,
+    *,
+    axis: Literal[0, 1],
+    dtype: DTypeLike | None = None,
 ) -> NDArray[Any]: ...
 @overload
 def sum(
@@ -47,7 +54,12 @@ def sum(
 
 
 def sum(
-    x: ArrayLike | CPUArray | types.DaskArray,
+    x: ArrayLike
+    | NDArray[Any]
+    | types.CSBase
+    | types.H5Dataset
+    | types.ZarrArray
+    | types.DaskArray,
     /,
     *,
     axis: Literal[0, 1, None] = None,
@@ -71,7 +83,12 @@ def sum(
 
 @singledispatch
 def _sum(
-    x: ArrayLike | CPUArray | types.DaskArray,
+    x: ArrayLike
+    | NDArray[Any]
+    | types.CSBase
+    | types.H5Dataset
+    | types.ZarrArray
+    | types.DaskArray,
     /,
     *,
     axis: Literal[0, 1, None] = None,

From 17fdffa51064da842f12a9dd2163180af6c74844 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Mon, 10 Mar 2025 15:19:20 +0100
Subject: [PATCH 15/21] only sparse are limited to float

---
 src/testing/fast_array_utils/_array_type.py |  5 +++++
 tests/test_stats.py                         | 14 ++++++++++----
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/testing/fast_array_utils/_array_type.py b/src/testing/fast_array_utils/_array_type.py
index cfd1f36..cef7715 100644
--- a/src/testing/fast_array_utils/_array_type.py
+++ b/src/testing/fast_array_utils/_array_type.py
@@ -145,6 +145,11 @@ def cls(self) -> type[Arr]:  # noqa: PLR0911
                 msg = f"Unknown array class: {self}"
                 raise ValueError(msg)
 
+    @cached_property
+    def classes(self) -> tuple[type[Array], ...]:
+        """Array classes for :func:`isinstance` checks (including the inner one for dask)."""
+        return (self.cls, *(() if self.inner is None else self.inner.classes))
+
     def random(
         self,
         shape: tuple[int, int],
diff --git a/tests/test_stats.py b/tests/test_stats.py
index 9801e84..2524d52 100644
--- a/tests/test_stats.py
+++ b/tests/test_stats.py
@@ -63,8 +63,11 @@ def test_sum(
     axis: Literal[0, 1, None],
 ) -> None:
     np_arr = np.array([[1, 2, 3], [4, 5, 6]], dtype=dtype_in)
-    if (array_type.flags & Flags.Gpu) and np_arr.dtype.kind != "f":
-        pytest.skip("GPU arrays only support floats")
+    if (
+        any(issubclass(cls, types.CupyCSMatrix) for cls in array_type.classes)
+        and np_arr.dtype.kind != "f"
+    ):
+        pytest.skip("CuPy sparse matrices only support floats")
     arr = array_type(np_arr.copy())
     assert arr.dtype == dtype_in
 
@@ -105,8 +108,11 @@ def test_mean(
     array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: float | list[float]
 ) -> None:
     np_arr = np.array([[1, 2, 3], [4, 5, 6]])
-    if (array_type.flags & Flags.Gpu) and np_arr.dtype.kind != "f":
-        pytest.skip("GPU arrays only support floats")
+    if (
+        any(issubclass(cls, types.CupyCSMatrix) for cls in array_type.classes)
+        and np_arr.dtype.kind != "f"
+    ):
+        pytest.skip("CuPy sparse matrices only support floats")
     np.testing.assert_array_equal(np.mean(np_arr, axis=axis), expected)
 
     arr = array_type(np_arr)

From 862faf4ea8657fa58a0d67af8172616d42ec9b21 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Mon, 10 Mar 2025 15:22:54 +0100
Subject: [PATCH 16/21] simpler

---
 tests/test_stats.py | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/tests/test_stats.py b/tests/test_stats.py
index 2524d52..7bb5ed0 100644
--- a/tests/test_stats.py
+++ b/tests/test_stats.py
@@ -63,10 +63,7 @@ def test_sum(
     axis: Literal[0, 1, None],
 ) -> None:
     np_arr = np.array([[1, 2, 3], [4, 5, 6]], dtype=dtype_in)
-    if (
-        any(issubclass(cls, types.CupyCSMatrix) for cls in array_type.classes)
-        and np_arr.dtype.kind != "f"
-    ):
+    if array_type in ATS_CUPY_SPARSE and np_arr.dtype.kind != "f":
         pytest.skip("CuPy sparse matrices only support floats")
     arr = array_type(np_arr.copy())
     assert arr.dtype == dtype_in
@@ -108,10 +105,7 @@ def test_mean(
     array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: float | list[float]
 ) -> None:
     np_arr = np.array([[1, 2, 3], [4, 5, 6]])
-    if (
-        any(issubclass(cls, types.CupyCSMatrix) for cls in array_type.classes)
-        and np_arr.dtype.kind != "f"
-    ):
+    if array_type in ATS_CUPY_SPARSE and np_arr.dtype.kind != "f":
         pytest.skip("CuPy sparse matrices only support floats")
     np.testing.assert_array_equal(np.mean(np_arr, axis=axis), expected)
 

From b2151a374dbe3e3c0ff9041393025940530421dc Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Tue, 11 Mar 2025 09:28:08 +0100
Subject: [PATCH 17/21]


From 9be7813fa9c71bfd71e790d31b321a1087111687 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Tue, 11 Mar 2025 09:43:21 +0100
Subject: [PATCH 18/21] no arraylike

---
 src/fast_array_utils/stats/__init__.py | 14 +++++---------
 src/fast_array_utils/stats/_sum.py     |  8 +++-----
 tests/test_to_dense.py                 |  4 +++-
 3 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/src/fast_array_utils/stats/__init__.py b/src/fast_array_utils/stats/__init__.py
index b46dde7..15203e1 100644
--- a/src/fast_array_utils/stats/__init__.py
+++ b/src/fast_array_utils/stats/__init__.py
@@ -17,7 +17,7 @@
     from typing import Any, Literal
 
     import numpy as np
-    from numpy.typing import ArrayLike, DTypeLike, NDArray
+    from numpy.typing import DTypeLike, NDArray
     from optype.numpy import ToDType
 
     from .. import types
@@ -185,15 +185,11 @@ def mean_var(
 # https://github.com/scverse/fast-array-utils/issues/52
 @overload
 def sum(
-    x: ArrayLike | CpuArray | GpuArray | DiskArray,
-    /,
-    *,
-    axis: None = None,
-    dtype: DTypeLike | None = None,
+    x: CpuArray | GpuArray | DiskArray, /, *, axis: None = None, dtype: DTypeLike | None = None
 ) -> np.number[Any]: ...
 @overload
 def sum(
-    x: ArrayLike | CpuArray | DiskArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None
+    x: CpuArray | DiskArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None
 ) -> NDArray[Any]: ...
 @overload
 def sum(
@@ -206,7 +202,7 @@ def sum(
 
 
 def sum(
-    x: ArrayLike | CpuArray | GpuArray | DiskArray | types.DaskArray,
+    x: CpuArray | GpuArray | DiskArray | types.DaskArray,
     /,
     *,
     axis: Literal[0, 1, None] = None,
@@ -225,4 +221,4 @@ def sum(
 
     """
     validate_axis(axis)
-    return sum_(x, axis=axis, dtype=dtype)  # type: ignore[arg-type]  # literally the same type, wtf mypy
+    return sum_(x, axis=axis, dtype=dtype)
diff --git a/src/fast_array_utils/stats/_sum.py b/src/fast_array_utils/stats/_sum.py
index 4d28d9e..feb64bf 100644
--- a/src/fast_array_utils/stats/_sum.py
+++ b/src/fast_array_utils/stats/_sum.py
@@ -12,7 +12,6 @@
 if TYPE_CHECKING:
     from typing import Any, Literal
 
-    from numpy._typing._array_like import _ArrayLikeFloat_co as ArrayLike
     from numpy.typing import DTypeLike, NDArray
 
     from ..typing import CpuArray, DiskArray, GpuArray
@@ -20,7 +19,7 @@
 
 @singledispatch
 def sum_(
-    x: ArrayLike | CpuArray | GpuArray | DiskArray | types.DaskArray,
+    x: CpuArray | GpuArray | DiskArray | types.DaskArray,
     /,
     *,
     axis: Literal[0, 1, None] = None,
@@ -52,9 +51,8 @@ def _sum_cs(
 
     if isinstance(x, types.CSMatrix):
         x = sp.csr_array(x) if x.format == "csr" else sp.csc_array(x)
-    if TYPE_CHECKING:
-        assert isinstance(x, ArrayLike)  # pyright: ignore[reportArgumentType]
-    return cast("NDArray[Any] | np.number[Any]", np.sum(x, axis=axis, dtype=dtype))
+
+    return cast("NDArray[Any] | np.number[Any]", np.sum(x, axis=axis, dtype=dtype))  # type: ignore[call-overload]
 
 
 @sum_.register(types.DaskArray)
diff --git a/tests/test_to_dense.py b/tests/test_to_dense.py
index 0930d22..77e9d23 100644
--- a/tests/test_to_dense.py
+++ b/tests/test_to_dense.py
@@ -11,10 +11,12 @@
 
 
 if TYPE_CHECKING:
+    from typing import TypeAlias
+
     from fast_array_utils.typing import CpuArray, DiskArray, GpuArray
     from testing.fast_array_utils import ArrayType
 
-    Array = CpuArray | GpuArray | DiskArray | types.CSDataset | types.DaskArray
+    Array: TypeAlias = CpuArray | GpuArray | DiskArray | types.CSDataset | types.DaskArray
 
 
 @pytest.mark.parametrize("to_memory", [True, False], ids=["to_memory", "not_to_memory"])

From d45416ba471fc177385350de0345733f9a848255 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Tue, 11 Mar 2025 10:23:49 +0100
Subject: [PATCH 19/21] simplify

---
 src/fast_array_utils/types.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/src/fast_array_utils/types.py b/src/fast_array_utils/types.py
index e70d43d..7759d4b 100644
--- a/src/fast_array_utils/types.py
+++ b/src/fast_array_utils/types.py
@@ -50,17 +50,13 @@
 """A sparse compressed matrix or array."""
 
 
-if TYPE_CHECKING or find_spec("cupy"):
+if TYPE_CHECKING or find_spec("cupy"):  # cupy always comes with cupyx
     from cupy import ndarray as CupyArray
-else:  # pragma: no cover
-    CupyArray = type("ndarray", (), {})
-    CupyArray.__module__ = "cupy"
-
-
-if TYPE_CHECKING or find_spec("cupyx"):
     from cupyx.scipy.sparse import csc_matrix as CupyCSCMatrix
     from cupyx.scipy.sparse import csr_matrix as CupyCSRMatrix
 else:  # pragma: no cover
+    CupyArray = type("ndarray", (), {})
+    CupyArray.__module__ = "cupy"
     CupyCSCMatrix = type("csc_matrix", (), {})
     CupyCSRMatrix = type("csr_matrix", (), {})
     CupyCSCMatrix.__module__ = CupyCSRMatrix.__module__ = "cupyx.scipy.sparse"

From 5978d1ed7a586bfeba23d24b1db0e5580fcd607d Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Tue, 11 Mar 2025 10:41:27 +0100
Subject: [PATCH 20/21] codecov

---
 .github/workflows/ci-gpu.yml | 30 ++++++++++++++++--------------
 .github/workflows/ci.yml     |  7 +++++--
 2 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/.github/workflows/ci-gpu.yml b/.github/workflows/ci-gpu.yml
index 402eaf1..08dbb54 100644
--- a/.github/workflows/ci-gpu.yml
+++ b/.github/workflows/ci-gpu.yml
@@ -37,32 +37,34 @@ jobs:
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0
-
       - name: Nvidia SMI sanity check
         run: nvidia-smi
-
-      - name: Install Python
-        uses: actions/setup-python@v5
+      - uses: actions/setup-python@v5
         with:
           python-version: "3.12"
-
-      - name: Install UV
-        uses: hynek/setup-cached-uv@v2
+      - uses: hynek/setup-cached-uv@v2
         with:
           cache-dependency-path: pyproject.toml
-
       - name: Install package
         run: uv pip install --system -e .[test,full] cupy-cuda12x --extra-index-url=https://pypi.nvidia.com --index-strategy=unsafe-best-match
-
       - name: List installed packages
         run: uv pip list
-
       - name: Run tests
-        run: pytest
-
-      - name: Remove 'run-gpu-ci' Label
+        run: |
+          coverage run -m pytest -m "not benchmark"
+          coverage report
+          # https://github.com/codecov/codecov-cli/issues/648
+          coverage xml
+          rm test-data/.coverage
+      - uses: codecov/codecov-action@v5
+        with:
+          name: GPU tests
+          fail_ci_if_error: true
+          files: test-data/coverage.xml
+          token: ${{ secrets.CODECOV_TOKEN }}
+      - name: Remove “run-gpu-ci” Label
         if: always()
         uses: actions-ecosystem/action-remove-labels@v1
         with:
-          labels: 'run-gpu-ci'
+          labels: run-gpu-ci
           github_token: ${{ secrets.GITHUB_TOKEN }}
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 94a7493..5154c0f 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -13,11 +13,11 @@ env:
 
 jobs:
   test:
+    name: Minimal test job
     runs-on: ubuntu-latest
     strategy:
       matrix:
         python-version: ["3.11", "3.13"]
-        extras: [min, full]
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
@@ -27,7 +27,7 @@ jobs:
         with:
           enable-cache: true
           cache-dependency-glob: pyproject.toml
-      - run: uv pip install --system -e .[test${{ matrix.extras == 'full' && ',full' || '' }}]
+      - run: uv pip install --system -e .[test]
       - run: |
           coverage run -m pytest -m "not benchmark"
           coverage report
@@ -36,10 +36,12 @@ jobs:
           rm test-data/.coverage
       - uses: codecov/codecov-action@v5
         with:
+          name: Min tests
           fail_ci_if_error: true
           files: test-data/coverage.xml
           token: ${{ secrets.CODECOV_TOKEN }}
   bench:
+    name: CPU Benchmarks
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -56,6 +58,7 @@ jobs:
           run: pytest -m benchmark --codspeed
           token: ${{ secrets.CODSPEED_TOKEN }}
   check:
+    name: Static checks
     runs-on: ubuntu-latest
     strategy:
       matrix:

From 39a9da2d28f667838122eecf59ae2af43f6f0605 Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Tue, 11 Mar 2025 10:56:48 +0100
Subject: [PATCH 21/21] renames

---
 .github/workflows/ci-gpu.yml | 7 ++++---
 .github/workflows/ci.yml     | 8 ++++----
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/ci-gpu.yml b/.github/workflows/ci-gpu.yml
index 08dbb54..b1577ed 100644
--- a/.github/workflows/ci-gpu.yml
+++ b/.github/workflows/ci-gpu.yml
@@ -20,13 +20,14 @@ concurrency:
 
 jobs:
   check:
+    name: Check Label
     runs-on: ubuntu-latest
     steps:
       - uses: flying-sheep/check@v1
         with:
           success: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'run-gpu-ci') }}
   test:
-    name: GPU Tests
+    name: All Tests
     needs: check
     runs-on: "cirun-aws-gpu--${{ github.run_id }}"
     timeout-minutes: 30
@@ -37,7 +38,7 @@ jobs:
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0
-      - name: Nvidia SMI sanity check
+      - name: Check NVIDIA SMI
         run: nvidia-smi
       - uses: actions/setup-python@v5
         with:
@@ -58,7 +59,7 @@ jobs:
           rm test-data/.coverage
       - uses: codecov/codecov-action@v5
         with:
-          name: GPU tests
+          name: GPU Tests
           fail_ci_if_error: true
           files: test-data/coverage.xml
           token: ${{ secrets.CODECOV_TOKEN }}
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 5154c0f..944d746 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -1,6 +1,6 @@
 # https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
 
-name: Python
+name: CI
 
 on:
   push:
@@ -13,7 +13,7 @@ env:
 
 jobs:
   test:
-    name: Minimal test job
+    name: Min Tests
     runs-on: ubuntu-latest
     strategy:
       matrix:
@@ -36,7 +36,7 @@ jobs:
           rm test-data/.coverage
       - uses: codecov/codecov-action@v5
         with:
-          name: Min tests
+          name: Min Tests
           fail_ci_if_error: true
           files: test-data/coverage.xml
           token: ${{ secrets.CODECOV_TOKEN }}
@@ -58,7 +58,7 @@ jobs:
           run: pytest -m benchmark --codspeed
           token: ${{ secrets.CODSPEED_TOKEN }}
   check:
-    name: Static checks
+    name: Static Checks
     runs-on: ubuntu-latest
     strategy:
       matrix: