import fix

Author: amalia-k510

src/fast_array_utils/stats/_generic_ops.py

@@ -4,7 +4,6 @@
 from functools import singledispatch
 from typing import TYPE_CHECKING, cast, get_args
 
-import array_api_compat
 import numpy as np
 
 from .. import types
@@ -45,7 +44,7 @@ def generic_op(
 
 
 @generic_op.register(np.ndarray)
-# to avoid going array api path that would slow down the performance
+# register explicitly to avoid the array API path and performance slow down
 def _generic_op_numpy(
     x: np.ndarray,
     /,
@@ -73,6 +72,8 @@ def _generic_op_array_api(
     """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))
 

src/fast_array_utils/stats/_power.py

@@ -4,7 +4,6 @@
 from functools import singledispatch
 from typing import TYPE_CHECKING
 
-import array_api_compat
 import numpy as np
 
 from .. import types
@@ -40,6 +39,7 @@ def _power_numpy(x: np.ndarray, n: int, /, dtype: DTypeLike | None = None) -> np
 
 @_power.register(types.HasArrayNamespace)
 def _power_array_api(x: types.HasArrayNamespace, n: int, /, dtype: DTypeLike | None = None) -> types.HasArrayNamespace:
+    import array_api_compat
 
     xp = array_api_compat.array_namespace(x)
     return xp.pow(x, n) if dtype is None else xp.pow(xp.astype(x, dtype), n)