Addressed the comments

Author: amalia-k510

pyproject.toml

@@ -50,15 +50,6 @@ doc = [
   "sphinx>=9.0.1",
   "sphinx-autofixture>=0.4.1",
 ]
-# for update-mypy-hook
-mypy = [
-  "fast-array-utils[full]",
-  "scipy-stubs",
-  # TODO: replace sphinx with this: { include-group = "doc" },
-  "sphinx",
-  "types-docutils",
-  { include-group = "test" },
-]
 test-min = [
   "coverage[toml]",
   "fast-array-utils[sparse,testing]", # include sparse for testing numba-less to_dense
@@ -75,7 +66,7 @@ envs.docs.dependency-groups = [ "doc" ]
 envs.docs.scripts.build = "sphinx-build -M html docs docs/_build"
 envs.docs.scripts.clean = "git clean -fdX docs"
 envs.docs.scripts.open = "python -m webbrowser -t docs/_build/html/index.html"
-envs.hatch-test.default-args = []
+envs.hatch-test.default-args = [  ]
 envs.hatch-test.dependency-groups = [ "test-min" ]
 # TODO: remove scipy once https://github.com/pypa/hatch/pull/2127 is released
 envs.hatch-test.extra-dependencies = [ "ipykernel", "ipycytoscape", "scipy" ]
@@ -100,7 +91,7 @@ metadata.hooks.docstring-description = {}
 metadata.hooks.fancy-pypi-readme.content-type = "text/x-rst"
 metadata.hooks.fancy-pypi-readme.fragments = [ { path = "README.rst", start-after = ".. begin" } ]
 version.source = "vcs"
-version.raw-options = { local_scheme = "no-local-version" }  # be able to publish dev version
+version.raw-options = { local_scheme = "no-local-version" } # be able to publish dev version
 
 [tool.uv]
 override-dependencies = [ "sphinx>=9.0.1" ]
@@ -141,7 +132,7 @@ lint.per-file-ignores."typings/**/*.pyi" = [ "A002", "F403", "F405", "N801" ]
 lint.allowed-confusables = [ "×", "’" ]
 lint.flake8-bugbear.extend-immutable-calls = [ "testing.fast_array_utils.Flags" ]
 lint.flake8-copyright.notice-rgx = "SPDX-License-Identifier: MPL-2\\.0"
-lint.flake8-type-checking.exempt-modules = []
+lint.flake8-type-checking.exempt-modules = [  ]
 lint.flake8-type-checking.strict = true
 lint.isort.known-first-party = [ "fast_array_utils" ]
 lint.isort.lines-after-imports = 2

src/fast_array_utils/stats/_is_constant.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from functools import partial, singledispatch
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING
 
 import numba
 import numpy as np
@@ -15,8 +15,6 @@
 
     from numpy.typing import NDArray
 
-# checking if all values in an array are the same
-
 
 @singledispatch
 def is_constant_(
@@ -26,36 +24,30 @@ def is_constant_(
     axis: Literal[0, 1] | None = None,
 ) -> Any:  # noqa: ANN401
 
+    # Catch types that lack __array_namespace like PyTorch
     import array_api_compat
 
     if array_api_compat.is_array_api_obj(a):
-        xp = array_api_compat.array_namespace(a)
-        match axis:
-            case None:
-                return bool((a == xp.reshape(a, (-1,))[0]).all())
-            case 0:
-                return is_constant_(a.T, axis=1)  # reusing axis = 1
-            case 1:
-                b = xp.broadcast_to(a[:, 0:1], a.shape)
-                return (a == b).all(axis=1)
+        return _is_constant_ndarray(a, axis=axis)
     raise NotImplementedError
 
 
-@is_constant_.register(np.ndarray | types.CupyArray)
+@is_constant_.register(np.ndarray | types.CupyArray | types.HasArrayNamespace)
 def _is_constant_ndarray(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:
-            return bool((a == a.flat[0]).all())
+            return bool((a == a.reshape(-1)[0]).all())
         case 0:
             return _is_constant_rows(a.T)
         case 1:
             return _is_constant_rows(a)
 
 
 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))
+    # broadcasts without needing np.broadcast_to
+    return (a == a[:, 0:1]).all(axis=1)
 
 
 @is_constant_.register(types.CSBase)

src/fast_array_utils/stats/_mean_var.py

@@ -45,8 +45,8 @@ def mean_var_(
     if axis is not None and isinstance(x, types.CSBase):
         mean_, var = _sparse_mean_var(x, axis=axis)
     else:
-        mean_ = mean(x, axis=axis, dtype=float64)
-        mean_sq = mean(power(x, 2, dtype=float64), axis=axis) if isinstance(x, types.DaskArray) else mean(power(x, 2), axis=axis, dtype=float64)
+        mean_ = mean(x, axis=axis, dtype=xp.float64)
+        mean_sq = mean(power(x, 2, dtype=xp.float64), axis=axis) if isinstance(x, types.DaskArray) else mean(power(x, 2), axis=axis, dtype=xp.float64)
         var = mean_sq - mean_**2
     if correction:  # R convention == 1 (unbiased estimator)
         n = np.prod(x.shape) if axis is None else x.shape[axis]