CuPy support (#51)

Author: flying-sheep

.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

.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

.github/workflows/ci-gpu.yml

@@ -0,0 +1,71 @@
+name: GPU-CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    types:
+      - labeled
+      - 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
+
+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: All 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: Check NVIDIA SMI
+        run: nvidia-smi
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - 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: |
+          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
+          github_token: ${{ secrets.GITHUB_TOKEN }}

.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,11 +13,11 @@ env:
 
 jobs:
   test:
+    name: Min Tests
     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:

src/fast_array_utils/conv/_to_dense.py

@@ -37,11 +37,9 @@ 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
-
     from . import to_dense
 
-    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]
 
 
@@ -56,7 +54,7 @@ def _to_dense_ooc(x: types.CSDataset, /, *, to_memory: bool = False) -> NDArray[
     return to_dense(cast("types.CSBase", x.to_memory()))
 
 
-@to_dense_.register(GpuArray)  # type: ignore[call-overload,misc]
+@to_dense_.register(types.CupyArray | types.CupyCSMatrix)  # type: ignore[call-overload,misc]
 def _to_dense_cupy(x: GpuArray, /, *, 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

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
@@ -27,16 +27,23 @@
 
 
 @overload
-def is_constant(a: types.DaskArray, /, *, axis: Literal[0, 1, None] = None) -> types.DaskArray: ...
+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: CpuArray, /, *, axis: None = None) -> bool: ...
+def is_constant(a: types.CupyArray, /, *, axis: Literal[0, 1]) -> types.CupyArray: ...
 @overload
-def is_constant(a: CpuArray, /, *, axis: Literal[0, 1]) -> NDArray[np.bool]: ...
+def is_constant(a: types.DaskArray, /, *, axis: Literal[0, 1, None] = None) -> types.DaskArray: ...
 
 
 def is_constant(
-    a: CpuArray | 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
@@ -82,9 +89,13 @@ def mean(
 ) -> np.number[Any]: ...
 @overload
 def mean(
-    x: CpuArray | GpuArray | DiskArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None
+    x: CpuArray | DiskArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None
 ) -> NDArray[np.number[Any]]: ...
 @overload
+def mean(
+    x: GpuArray, /, *, 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: ...
@@ -96,7 +107,7 @@ def mean(
     *,
     axis: Literal[0, 1, None] = None,
     dtype: DTypeLike | None = None,
-) -> NDArray[np.number[Any]] | np.number[Any] | types.DaskArray:
+) -> NDArray[np.number[Any]] | types.CupyArray | np.number[Any] | types.DaskArray:
     """Mean over both or one axis.
 
     Returns
@@ -115,11 +126,15 @@ def mean(
 @overload
 def mean_var(
     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: CpuArray | GpuArray, /, *, 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
@@ -133,8 +148,9 @@ def mean_var(
     axis: Literal[0, 1, None] = None,
     correction: int = 0,
 ) -> (
-    tuple[NDArray[np.float64], NDArray[np.float64]]
-    | tuple[np.float64, np.float64]
+    tuple[np.float64, np.float64]
+    | tuple[NDArray[np.float64], NDArray[np.float64]]
+    | tuple[types.CupyArray, types.CupyArray]
     | tuple[types.DaskArray, types.DaskArray]
 ):
     """Mean and variance over both or one axis.
@@ -169,33 +185,29 @@ 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 | GpuArray | 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(
+    x: GpuArray, /, *, 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 | CpuArray | GpuArray | DiskArray | types.DaskArray,
+    x: CpuArray | GpuArray | DiskArray | types.DaskArray,
     /,
     *,
     axis: Literal[0, 1, None] = None,
     dtype: DTypeLike | None = None,
-) -> NDArray[Any] | np.number[Any] | types.DaskArray:
+) -> NDArray[Any] | types.CupyArray | np.number[Any] | types.DaskArray:
     """Sum over both or one axis.
 
     Returns
@@ -209,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)

src/fast_array_utils/stats/_is_constant.py

@@ -20,15 +20,18 @@
 
 @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
 
 
-@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:
@@ -39,7 +42,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))
 

src/fast_array_utils/stats/_mean_var.py

@@ -27,6 +27,7 @@ def mean_var_(
     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]
 ):

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)