types

Author: flying-sheep

.pre-commit-config.yaml

@@ -32,6 +32,7 @@ repos:
     - array-api-compat>=1.13
     - dask>=2026.1
     - h5py>=3.15
+    - jax>=0.10
     - numba>=0.63
     - packaging>=26
     - pytest>=9

src/fast_array_utils/conv/_to_dense.py

@@ -47,7 +47,7 @@ def _to_dense_numpy(x: np.ndarray, /, *, order: Literal["K", "A", "C", "F"] = "K
 
 
 @to_dense_.register(types.HasArrayNamespace)
-def _to_dense_array_api(x: types.HasArrayNamespace, /, *, order: Literal["K", "A", "C", "F"] = "K", to_cpu_memory: bool = False) -> Any:  # noqa: ANN401
+def _to_dense_array_api[A: types.HasArrayNamespace](x: A, /, *, order: Literal["K", "A", "C", "F"] = "K", to_cpu_memory: bool = False) -> A | np.ndarray:
     if to_cpu_memory:
         return np.asarray(x, order=order)
     return x

src/fast_array_utils/stats/__init__.py

@@ -214,13 +214,13 @@ def _mk_generic_op(op: DtypeOps) -> StatFunDtype: ...
 # https://github.com/scverse/fast-array-utils/issues/52
 def _mk_generic_op(op: Ops) -> StatFunNoDtype | StatFunDtype:
     def _generic_op(
-        x: CpuArray | GpuArray | DiskArray | types.DaskArray,
+        x: CpuArray | GpuArray | DiskArray | types.DaskArray | types.HasArrayNamespace,
         /,
         *,
         axis: Literal[0, 1] | None = None,
         dtype: DTypeLike | None = None,
         keep_cupy_as_array: bool = False,
-    ) -> NDArray[Any] | np.number[Any] | types.CupyArray | types.DaskArray:
+    ) -> NDArray[Any] | np.number[Any] | types.CupyArray | types.DaskArray | types.HasArrayNamespace:
         from ._generic_ops import generic_op
 
         assert dtype is None or op in get_args(DtypeOps), f"`dtype` is not supported for operation {op!r}"
@@ -359,14 +359,16 @@ def sum(x: GpuArray, /, *, axis: None = None, dtype: DTypeLike | None = None, ke
 def sum(x: GpuArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None, keep_cupy_as_array: bool = False) -> types.CupyArray: ...
 @overload
 def sum(x: types.DaskArray, /, *, axis: Literal[0, 1] | None = None, dtype: DTypeLike | None = None, keep_cupy_as_array: bool = False) -> types.DaskArray: ...
+@overload
+def sum[A: types.HasArrayNamespace](x: A, /, *, axis: Literal[0, 1] | None = None, dtype: DTypeLike | None = None, keep_cupy_as_array: bool = False) -> A: ...
 def sum(
-    x: CpuArray | GpuArray | DiskArray | types.DaskArray,
+    x: CpuArray | GpuArray | DiskArray | types.DaskArray | types.HasArrayNamespace,
     /,
     *,
     axis: Literal[0, 1] | None = None,
     dtype: DTypeLike | None = None,
     keep_cupy_as_array: bool = False,
-) -> NDArray[Any] | types.CupyArray | np.number[Any] | types.DaskArray:
+) -> NDArray[Any] | types.CupyArray | np.number[Any] | types.DaskArray | types.HasArrayNamespace:
     """Sum over both or one axis.
 
     Parameters

src/fast_array_utils/stats/_generic_ops.py

@@ -59,22 +59,22 @@ def _generic_op_numpy(
 
 
 @generic_op.register(types.HasArrayNamespace)
-def _generic_op_array_api(
-    x: types.HasArrayNamespace,
+def _generic_op_array_api[A: types.HasArrayNamespace](
+    x: A,
     /,
     op: Ops,
     *,
     axis: Literal[0, 1] | None = None,
     dtype: DTypeLike | None = None,
     keep_cupy_as_array: bool = False,
-) -> Any:  # noqa: ANN401
+) -> A:
     """Handle arrays with native array API support."""
     del keep_cupy_as_array
 
     import array_api_compat
 
     xp = array_api_compat.array_namespace(x)
-    return getattr(xp, op)(x, axis=axis, **_dtype_kw(dtype, op))
+    return getattr(xp, op)(x, axis=axis, **_dtype_kw(dtype, op))  # type: ignore[no-any-return]
 
 
 @generic_op.register(types.CupyArray | types.CupyCSMatrix)

tests/test_jax.py

@@ -12,7 +12,7 @@
 
 
 if TYPE_CHECKING:
-    from typing import Any, Literal
+    from typing import Literal
 
 pytestmark = pytest.mark.skipif(not find_spec("jax"), reason="jax not installed")
 
@@ -21,18 +21,18 @@
     # problem as mean_var passes dtype= np.float64 internally, which crashes without this fix
     import jax
 
-    jax.config.update("jax_enable_x64", True)  # noqa: FBT003
+    jax.config.update("jax_enable_x64", True)  # type: ignore[no-untyped-call]  # noqa: FBT003
 
 
 @pytest.fixture
-def jax_arr() -> Any:  # noqa: ANN401
+def jax_arr() -> jax.Array:
     import jax.numpy as jnp
 
     return jnp.array([[1, 0], [2, 0], [3, 0]], dtype=jnp.float32)
 
 
 @pytest.mark.parametrize("axis", [None, 0, 1])
-def test_sum(jax_arr: Any, axis: Literal[0, 1] | None) -> None:  # noqa: ANN401
+def test_sum(jax_arr: jax.Array, axis: Literal[0, 1] | None) -> None:
     import jax.numpy as jnp
 
     result = stats.sum(jax_arr, axis=axis)
@@ -41,7 +41,7 @@ def test_sum(jax_arr: Any, axis: Literal[0, 1] | None) -> None:  # noqa: ANN401
 
 
 @pytest.mark.parametrize("axis", [None, 0, 1])
-def test_min(jax_arr: Any, axis: Literal[0, 1] | None) -> None:  # noqa: ANN401
+def test_min(jax_arr: jax.Array, axis: Literal[0, 1] | None) -> None:
     import jax.numpy as jnp
 
     result = stats.min(jax_arr, axis=axis)
@@ -50,7 +50,7 @@ def test_min(jax_arr: Any, axis: Literal[0, 1] | None) -> None:  # noqa: ANN401
 
 
 @pytest.mark.parametrize("axis", [None, 0, 1])
-def test_max(jax_arr: Any, axis: Literal[0, 1] | None) -> None:  # noqa: ANN401
+def test_max(jax_arr: jax.Array, axis: Literal[0, 1] | None) -> None:
     import jax.numpy as jnp
 
     result = stats.max(jax_arr, axis=axis)
@@ -59,7 +59,7 @@ def test_max(jax_arr: Any, axis: Literal[0, 1] | None) -> None:  # noqa: ANN401
 
 
 @pytest.mark.parametrize("axis", [None, 0, 1])
-def test_mean(jax_arr: Any, axis: Literal[0, 1] | None) -> None:  # noqa: ANN401
+def test_mean(jax_arr: jax.Array, axis: Literal[0, 1] | None) -> None:
     import jax.numpy as jnp
 
     result = stats.mean(jax_arr, axis=axis)
@@ -95,7 +95,7 @@ def test_is_constant(axis: Literal[0, 1] | None) -> None:
 
 
 @pytest.mark.parametrize("axis", [None, 0, 1])
-def test_mean_var(jax_arr: Any, axis: Literal[0, 1] | None) -> None:  # noqa: ANN401
+def test_mean_var(jax_arr: jax.Array, axis: Literal[0, 1] | None) -> None:
     import jax.numpy as jnp
 
     mean, var = stats.mean_var(jax_arr, axis=axis, correction=1)
@@ -108,14 +108,14 @@ def test_mean_var(jax_arr: Any, axis: Literal[0, 1] | None) -> None:  # noqa: AN
     assert jnp.allclose(var, var_expected)
 
 
-def test_to_dense(jax_arr: Any) -> None:  # noqa: ANN401
+def test_to_dense(jax_arr: jax.Array) -> None:
     import jax.numpy as jnp
 
     result = to_dense(jax_arr)
     assert jnp.array_equal(result, jax_arr)
 
 
-def test_to_dense_to_cpu(jax_arr: Any) -> None:  # noqa: ANN401
+def test_to_dense_to_cpu(jax_arr: jax.Array) -> None:
     result = to_dense(jax_arr, to_cpu_memory=True)
     assert isinstance(result, np.ndarray)
     np.testing.assert_array_equal(result, np.asarray(jax_arr))