diff --git a/src/fast_array_utils/stats/_mean_var.py b/src/fast_array_utils/stats/_mean_var.py index c704bc9..2b8d507 100644 --- a/src/fast_array_utils/stats/_mean_var.py +++ b/src/fast_array_utils/stats/_mean_var.py @@ -72,10 +72,10 @@ def _sparse_mean_var(mtx: types.CSBase, /, *, axis: Literal[0, 1]) -> tuple[NDAr elif mtx.format == "csc": ax_minor = 0 shape = mtx.shape[::-1] - else: + else: # pragma: no cover msg = "This function only works on sparse csr and csc matrices" raise TypeError(msg) - if len(shape) == 1: + if len(shape) == 1: # pragma: no cover msg = "array must have 2 dimensions" raise TypeError(msg) f = sparse_mean_var_major_axis if axis == ax_minor else sparse_mean_var_minor_axis diff --git a/src/fast_array_utils/types.py b/src/fast_array_utils/types.py index 6042540..38df35b 100644 --- a/src/fast_array_utils/types.py +++ b/src/fast_array_utils/types.py @@ -8,7 +8,9 @@ if TYPE_CHECKING: - from types import ModuleType + from typing import Self + + from array_api.latest import ArrayNamespace __all__ = [ @@ -139,5 +141,5 @@ def shape(self) -> tuple[int, ...]: def dtype(self) -> object: """The data type of the array.""" - def __array_namespace__(self, /, *, api_version: str | None = None) -> ModuleType: + def __array_namespace__(self, /, *, api_version: str | None = None) -> ArrayNamespace[Self, object, object]: # type: ignore[type-var] """Get Array API namespace.""" diff --git a/tests/test_stats.py b/tests/test_stats.py index 993020a..0860334 100644 --- a/tests/test_stats.py +++ b/tests/test_stats.py @@ -396,12 +396,16 @@ def test_stats_benchmark( n_elems, density = 10_000_000, 0.01 n = int(np.sqrt(n_elems / density if "sparse" in array_type.mod else n_elems)) arr = array_type.random((n, n), density=density, dtype=dtype) + if ( + func is stats.is_constant + and isinstance(arr, types.CSBase) + and ((array_type.name == "csr_array" and axis == 1) or (array_type.name == "csc_array" and axis == 0)) + ): + # random rows break at the first entry, so make every row constant (explicit zeros) to scan all entries + arr.data[:] = 0 func(arr, axis=axis) # warmup: numba compile - is_very_fast = func is stats.is_constant and ((array_type.name == "csr_array" and axis == 1) or (array_type.name == "csc_array" and axis == 0)) - @benchmark def call() -> None: - for _ in range(100 if is_very_fast else 1): - func(arr, axis=axis) + func(arr, axis=axis)