Reverts and tweaks

Author: crusaderky

array_api_compat/common/_aliases.py

@@ -6,7 +6,6 @@
 
 import inspect
 from collections.abc import Sequence
-from types import NoneType
 from typing import TYPE_CHECKING, Any, NamedTuple, cast
 
 from ._helpers import _check_device, array_namespace
@@ -384,7 +383,7 @@ def clip(
     out: Array | None = None,
 ) -> Array:
     def _isscalar(a: object) -> TypeIs[float | None]:
-        return isinstance(a, int | float | NoneType)
+        return isinstance(a, int | float) or a is None
 
     min_shape = () if _isscalar(min) else min.shape
     max_shape = () if _isscalar(max) else max.shape

array_api_compat/common/_helpers.py

@@ -12,14 +12,14 @@
 import math
 import sys
 import warnings
-from collections.abc import Hashable
+from collections.abc import Collection, Hashable
 from functools import lru_cache
-from types import NoneType
 from typing import (
     TYPE_CHECKING,
     Any,
     Final,
     Literal,
+    SupportsIndex,
     TypeAlias,
     TypeGuard,
     cast,
@@ -51,7 +51,7 @@
         | ndx.Array
         | sparse.SparseArray
         | torch.Tensor
-        | SupportsArrayNamespace
+        | SupportsArrayNamespace[Any]
     )
 
 _API_VERSIONS_OLD: Final = frozenset({"2021.12", "2022.12", "2023.12"})
@@ -630,9 +630,9 @@ def your_function(x, y):
                 raise ValueError(
                     "The given array does not have an array-api-compat wrapper"
                 )
-            x = cast(SupportsArrayNamespace, x)
+            x = cast("SupportsArrayNamespace[Any]", x)
             namespaces.add(x.__array_namespace__(api_version=api_version))
-        elif isinstance(x, int | float | complex | NoneType):
+        elif isinstance(x, int | float | complex) or x is None:
             continue
         else:
             # TODO: Support Python scalars?
@@ -890,12 +890,10 @@ def to_device(x: Array, device: Device, /, *, stream: int | Any | None = None) -
 
 
 @overload
-def size(x: HasShape[int]) -> int: ...
+def size(x: HasShape[Collection[SupportsIndex]]) -> int: ...
 @overload
-def size(x: HasShape[int | None]) -> int | None: ...
-@overload
-def size(x: HasShape[float]) -> int | None: ...  # Dask special case
-def size(x: HasShape[float | None]) -> int | None:
+def size(x: HasShape[Collection[SupportsIndex | None]]) -> int | None: ...
+def size(x: HasShape[Collection[SupportsIndex | None]]) -> int | None:
     """
     Return the total number of elements of x.
 
@@ -910,9 +908,9 @@ def size(x: HasShape[float | None]) -> int | None:
     # Lazy API compliant arrays, such as ndonnx, can contain None in their shape
     if None in x.shape:
         return None
-    out = math.prod(cast(tuple[float, ...], x.shape))
+    out = math.prod(cast("Collection[SupportsIndex]", x.shape))
     # dask.array.Array.shape can contain NaN
-    return None if math.isnan(out) else cast(int, out)
+    return None if math.isnan(out) else out
 
 
 @lru_cache(100)
@@ -1003,7 +1001,7 @@ def is_lazy_array(x: object) -> TypeGuard[_ArrayApiObj]:
     # on __bool__ (dask is one such example, which however is special-cased above).
 
     # Select a single point of the array
-    s = size(cast(HasShape, x))
+    s = size(cast("HasShape[Collection[SupportsIndex | None]]", x))
     if s is None:
         return True
     xp = array_namespace(x)

array_api_compat/common/_linalg.py

@@ -187,14 +187,14 @@ def vector_norm(
         # We can't reuse xp.linalg.norm(keepdims) because of the reshape hacks
         # above to avoid matrix norm logic.
         shape = list(x.shape)
-        _axis = cast(
+        axes = cast(
             "tuple[int, ...]",
             normalize_axis_tuple(  # pyright: ignore[reportCallIssue]
                 range(x.ndim) if axis is None else axis,
                 x.ndim,
             ),
         )
-        for i in _axis:
+        for i in axes:
             shape[i] = 1
         res = xp.reshape(res, tuple(shape))
 

array_api_compat/common/_typing.py

@@ -61,13 +61,13 @@ def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
     def __len__(self, /) -> int: ...
 
 
-class SupportsArrayNamespace(Protocol):
-    def __array_namespace__(self, /, *, api_version: str | None) -> Namespace: ...
+class SupportsArrayNamespace(Protocol[_T_co]):
+    def __array_namespace__(self, /, *, api_version: str | None) -> _T_co: ...
 
 
 class HasShape(Protocol[_T_co]):
     @property
-    def shape(self, /) -> tuple[_T_co, ...]: ...
+    def shape(self, /) -> _T_co: ...
 
 
 # Return type of `__array_namespace_info__.default_dtypes`

array_api_compat/cupy/_aliases.py

@@ -3,6 +3,7 @@
 from builtins import bool as py_bool
 
 import cupy as cp
+
 from ..common import _aliases, _helpers
 from ..common._typing import NestedSequence, SupportsBufferProtocol
 from .._internal import get_xp
@@ -119,7 +120,7 @@ def count_nonzero(
 
 
 # take_along_axis: axis defaults to -1 but in cupy (and numpy) axis is a required arg
-def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1):
+def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array:
     return cp.take_along_axis(x, indices, axis=axis)
 
 

array_api_compat/numpy/_aliases.py

@@ -119,14 +119,14 @@ def count_nonzero(
 ) -> Array:
     # NOTE: this is currently incorrectly typed in numpy, but will be fixed in
     # numpy 2.2.5 and 2.3.0: https://github.com/numpy/numpy/pull/28750
-    result = cast(Any, np.count_nonzero(x, axis=axis, keepdims=keepdims))  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType, reportCallIssue]
+    result = cast("Any", np.count_nonzero(x, axis=axis, keepdims=keepdims))  # pyright: ignore[reportArgumentType, reportCallIssue]
     if axis is None and not keepdims:
         return np.asarray(result)
     return result
 
 
 # take_along_axis: axis defaults to -1 but in numpy axis is a required arg
-def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1):
+def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array:
     return np.take_along_axis(x, indices, axis=axis)
 
 

array_api_compat/numpy/_info.py

@@ -7,7 +7,6 @@
 more details.
 
 """
-
 from __future__ import annotations
 
 from numpy import bool_ as bool
@@ -64,7 +63,7 @@ class __array_namespace_info__:
 
     """
 
-    __module__ = "numpy"
+    __module__ = 'numpy'
 
     def capabilities(self):
         """
@@ -183,7 +182,8 @@ def default_dtypes(
         """
         if device not in ["cpu", None]:
             raise ValueError(
-                f'Device not understood. Only "cpu" is allowed, but received: {device}'
+                'Device not understood. Only "cpu" is allowed, but received:'
+                f' {device}'
             )
         return {
             "real floating": dtype(float64),
@@ -254,7 +254,8 @@ def dtypes(
         """
         if device not in ["cpu", None]:
             raise ValueError(
-                f'Device not understood. Only "cpu" is allowed, but received: {device}'
+                'Device not understood. Only "cpu" is allowed, but received:'
+                f' {device}'
             )
         if kind is None:
             return {

array_api_compat/numpy/_typing.py

@@ -7,6 +7,7 @@
 Device: TypeAlias = Literal["cpu"]
 
 if TYPE_CHECKING:
+
     # NumPy 1.x on Python 3.10 fails to parse np.dtype[]
     DType: TypeAlias = np.dtype[
         np.bool_