STY: apply review suggestions

Co-authored-by: lucascolley <lucas.colley8@gmail.com>

Author: jorenham

array_api_compat/common/_aliases.py

@@ -652,7 +652,7 @@ def sign(x: Array, /, xp: Namespace, **kwargs: object) -> Array:
     if isdtype(x.dtype, "complex floating", xp=xp):
         out = (x / xp.abs(x, **kwargs))[...]
         # sign(0) = 0 but the above formula would give nan
-        out[x == 0 + 0j] = 0 + 0j
+        out[x == 0j] = 0j
     else:
         out = xp.sign(x, **kwargs)
     # CuPy sign() does not propagate nans. See

array_api_compat/common/_helpers.py

@@ -16,6 +16,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Final,
     Literal,
     SupportsIndex,
     TypeAlias,
@@ -56,6 +57,9 @@
         | _CupyArray
     )
 
+_API_VERSIONS_OLD: Final = frozenset({"2021.12", "2022.12", "2023.12"})
+_API_VERSIONS: Final = _API_VERSIONS_OLD | frozenset({"2024.12"})
+
 
 def _is_jax_zero_gradient_array(x: object) -> TypeGuard[_ZeroGradientArray]:
     """Return True if `x` is a zero-gradient array.
@@ -477,16 +481,11 @@ def is_array_api_strict_namespace(xp: Namespace) -> bool:
 
 
 def _check_api_version(api_version: str | None) -> None:
-    if api_version in ["2021.12", "2022.12", "2023.12"]:
+    if api_version in _API_VERSIONS_OLD:
         warnings.warn(
             f"The {api_version} version of the array API specification was requested but the returned namespace is actually version 2024.12"
         )
-    elif api_version is not None and api_version not in [
-        "2021.12",
-        "2022.12",
-        "2023.12",
-        "2024.12",
-    ]:
+    elif api_version is not None and api_version not in _API_VERSIONS:
         raise ValueError(
             "Only the 2024.12 version of the array API specification is currently supported"
         )

array_api_compat/dask/array/linalg.py

@@ -4,7 +4,7 @@
 
 import dask.array as da
 
-# These functions are in both the main and linalg namespaces
+# The `matmul` and `tensordot` functions are in both the main and linalg namespaces
 from dask.array import matmul, outer, tensordot
 
 # Exports

array_api_compat/numpy/_aliases.py

@@ -15,6 +15,8 @@
 if TYPE_CHECKING:
     from typing_extensions import Buffer, TypeIs
 
+# The values of the `_CopyMode` enum can be either `False`, `True`, or `2`:
+# https://github.com/numpy/numpy/blob/5a8a6a79d9c2fff8f07dcab5d41e14f8508d673f/numpy/_globals.pyi#L7-L10
 _Copy: TypeAlias = py_bool | Literal[2] | np._CopyMode
 
 bool = np.bool_
@@ -130,7 +132,9 @@ def count_nonzero(
     axis: int | tuple[int, ...] | None = None,
     keepdims: py_bool = False,
 ) -> Array:
-    result = cast("Any", np.count_nonzero(x, axis=axis, keepdims=keepdims))  # pyright: ignore
+    # 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))  # pyright: ignore[reportArgumentType, reportCallIssue]
     if axis is None and not keepdims:
         return np.asarray(result)
     return result