Repository: data-apis/array-api-strict
Pull request #135: TYP: type annotations
Author: login: crusaderky
Labels: name: Static Typing
Created: 2025-03-21T15:35:56Z | Merged: 2025-03-31T22:32:22Z

- Implement thorough static type annotations
- Align type annotations conventions to `array-api-compat` (https://github.com/data-apis/array-api-compat/pull/257)
- Add `py.typed`
- Manually run mypy and fix _most_ typing issues
- Manually run isort (one-off)
- Fix bug where `iinfo` and `finfo` would return a numpy dtype instead of a DType

# Out of scope
- Fix _all_ mypy errors (some highlight material issues that deserve their own PRs)
- Add mypy to CI
- Make pyright happy
- tests module


Diff:
diff --git a/array_api_strict/__init__.py b/array_api_strict/__init__.py
index e6a1763..116df25 100644
--- a/array_api_strict/__init__.py
+++ b/array_api_strict/__init__.py
@@ -16,6 +16,8 @@
 
 """
 
+from types import ModuleType
+
 __all__ = []
 
 # Warning: __array_api_version__ could change globally with
@@ -325,12 +327,16 @@
     ArrayAPIStrictFlags,
 )
 
-__all__ += ['set_array_api_strict_flags', 'get_array_api_strict_flags', 'reset_array_api_strict_flags', 'ArrayAPIStrictFlags']
+__all__ += [
+    'set_array_api_strict_flags',
+    'get_array_api_strict_flags',
+    'reset_array_api_strict_flags',
+    'ArrayAPIStrictFlags',
+    '__version__',
+]
 
 try:
-    from . import _version
-    __version__ = _version.__version__
-    del _version
+    from ._version import __version__  # type: ignore[import-not-found,unused-ignore]
 except ImportError:
     __version__ = "unknown"
 
@@ -340,7 +346,7 @@
 # use __getattr__. Note that linalg and fft are dynamically added and removed
 # from __all__ in set_array_api_strict_flags.
 
-def __getattr__(name):
+def __getattr__(name: str) -> ModuleType:
     if name in ['linalg', 'fft']:
         if name in get_array_api_strict_flags()['enabled_extensions']:
             if name == 'linalg':
diff --git a/array_api_strict/_array_object.py b/array_api_strict/_array_object.py
index 0595594..1304d5a 100644
--- a/array_api_strict/_array_object.py
+++ b/array_api_strict/_array_object.py
@@ -16,62 +16,70 @@
 from __future__ import annotations
 
 import operator
+import sys
+from collections.abc import Iterator
 from enum import IntEnum
+from types import ModuleType
+from typing import TYPE_CHECKING, Any, Final, Literal, SupportsIndex
 
-from ._creation_functions import asarray
+import numpy as np
+import numpy.typing as npt
+
+from ._creation_functions import Undef, _undef, asarray
 from ._dtypes import (
-    _DType,
+    DType,
     _all_dtypes,
     _boolean_dtypes,
+    _complex_floating_dtypes,
+    _dtype_categories,
+    _floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
-    _floating_dtypes,
-    _real_floating_dtypes,
-    _complex_floating_dtypes,
     _numeric_dtypes,
-    _result_type,
-    _dtype_categories,
+    _real_floating_dtypes,
     _real_to_complex_map,
+    _result_type,
 )
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
+from ._typing import PyCapsule
 
-from typing import TYPE_CHECKING, SupportsIndex
-import types
+if sys.version_info >= (3, 10):
+    from types import EllipsisType
+elif TYPE_CHECKING:
+    from typing_extensions import EllipsisType
+else:
+    EllipsisType = type(Ellipsis)
 
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union, Any
-    from ._typing import PyCapsule, Dtype
-    import numpy.typing as npt
-
-import numpy as np
 
 class Device:
-    def __init__(self, device="CPU_DEVICE"):
+    _device: Final[str]
+    __slots__ = ("_device", "__weakref__")
+
+    def __init__(self, device: str = "CPU_DEVICE"):
         if device not in ("CPU_DEVICE", "device1", "device2"):
             raise ValueError(f"The device '{device}' is not a valid choice.")
         self._device = device
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"array_api_strict.Device('{self._device}')"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if not isinstance(other, Device):
             return False
         return self._device == other._device
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         return hash(("Device", self._device))
 
 
 CPU_DEVICE = Device()
 ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"))
 
-_default = object()
-
 # See https://github.com/data-apis/array-api-strict/issues/67 and the comment
 # on __array__ below.
 _allow_array = True
 
+
 class Array:
     """
     n-d array object for the array API namespace.
@@ -87,12 +95,16 @@ class Array:
     functions, such as asarray().
 
     """
+
     _array: npt.NDArray[Any]
+    _dtype: DType
+    _device: Device
+    __slots__ = ("_array", "_dtype", "_device", "__weakref__")
 
     # Use a custom constructor instead of __init__, as manually initializing
     # this class is not supported API.
     @classmethod
-    def _new(cls, x, /, device):
+    def _new(cls, x: npt.NDArray[Any] | np.generic, /, device: Device | None) -> Array:
         """
         This is a private method for initializing the array API Array
         object.
@@ -107,7 +119,7 @@ def _new(cls, x, /, device):
         if isinstance(x, np.generic):
             # Convert the array scalar to a 0-D array
             x = np.asarray(x)
-        _dtype = _DType(x.dtype)
+        _dtype = DType(x.dtype)
         if _dtype not in _all_dtypes:
             raise TypeError(
                 f"The array_api_strict namespace does not support the dtype '{x.dtype}'"
@@ -120,7 +132,7 @@ def _new(cls, x, /, device):
         return obj
 
     # Prevent Array() from working
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *args: object, **kwargs: object) -> Array:
         raise TypeError(
             "The array_api_strict Array object should not be instantiated directly. Use an array creation function, such as asarray(), instead."
         )
@@ -128,7 +140,7 @@ def __new__(cls, *args, **kwargs):
     # These functions are not required by the spec, but are implemented for
     # the sake of usability.
 
-    def __repr__(self: Array, /) -> str:
+    def __repr__(self) -> str:
         """
         Performs the operation __repr__.
         """
@@ -159,7 +171,9 @@ def __repr__(self: Array, /) -> str:
     # This was implemented historically for compatibility, and removing it has
     # caused issues for some libraries (see
     # https://github.com/data-apis/array-api-strict/issues/67).
-    def __array__(self, dtype: None | np.dtype[Any] = None, copy: None | bool = None) -> npt.NDArray[Any]:
+    def __array__(
+        self, dtype: None | np.dtype[Any] = None, copy: None | bool = None
+    ) -> npt.NDArray[Any]:
         # We have to allow this to be internally enabled as there's no other
         # easy way to parse a list of Array objects in asarray().
         if _allow_array:
@@ -184,7 +198,9 @@ def __array__(self, dtype: None | np.dtype[Any] = None, copy: None | bool = None
     # spec in places where it either deviates from or is more strict than
     # NumPy behavior
 
-    def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_category: str, op: str) -> Array:
+    def _check_allowed_dtypes(
+        self, other: Array | bool | int | float | complex, dtype_category: str, op: str
+    ) -> Array:
         """
         Helper function for operators to only allow specific input dtypes
 
@@ -197,7 +213,7 @@ def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_categor
 
         if self.dtype not in _dtype_categories[dtype_category]:
             raise TypeError(f"Only {dtype_category} dtypes are allowed in {op}")
-        if isinstance(other, (int, complex, float, bool)):
+        if isinstance(other, (bool, int, float, complex)):
             other = self._promote_scalar(other)
         elif isinstance(other, Array):
             if other.dtype not in _dtype_categories[dtype_category]:
@@ -225,16 +241,18 @@ def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_categor
 
         return other
 
-    def _check_device(self, other):
+    def _check_device(self, other: Array | bool | int | float | complex) -> None:
         """Check that other is on a device compatible with the current array"""
-        if isinstance(other, (int, complex, float, bool)):
+        if isinstance(other, (bool, int, float, complex)):
             return
         elif isinstance(other, Array):
             if self.device != other.device:
                 raise ValueError(f"Arrays from two different devices ({self.device} and {other.device}) can not be combined.")
+        else:
+            raise TypeError(f"Expected Array | python scalar; got {type(other)}")
 
     # Helper function to match the type promotion rules in the spec
-    def _promote_scalar(self, scalar):
+    def _promote_scalar(self, scalar: bool | int | float | complex) -> Array:
         """
         Returns a promoted version of a Python scalar appropriate for use with
         operations on self.
@@ -291,7 +309,7 @@ def _promote_scalar(self, scalar):
         return Array._new(np.array(scalar, dtype=target_dtype._np_dtype), device=self.device)
 
     @staticmethod
-    def _normalize_two_args(x1, x2) -> Tuple[Array, Array]:
+    def _normalize_two_args(x1: Array, x2: Array) -> tuple[Array, Array]:
         """
         Normalize inputs to two arg functions to fix type promotion rules
 
@@ -327,7 +345,17 @@ def _normalize_two_args(x1, x2) -> Tuple[Array, Array]:
 
     # Note: A large fraction of allowed indices are disallowed here (see the
     # docstring below)
-    def _validate_index(self, key, op="getitem"):
+    def _validate_index(
+        self,
+        key: (
+            int
+            | slice
+            | EllipsisType
+            | Array
+            | tuple[int | slice | EllipsisType | Array | None, ...]
+        ),
+        op: Literal["getitem", "setitem"] = "getitem",
+    ) -> None:
         """
         Validate an index according to the array API.
 
@@ -509,7 +537,7 @@ def _validate_index(self, key, op="getitem"):
 
     # Everything below this line is required by the spec.
 
-    def __abs__(self: Array, /) -> Array:
+    def __abs__(self) -> Array:
         """
         Performs the operation __abs__.
         """
@@ -518,7 +546,7 @@ def __abs__(self: Array, /) -> Array:
         res = self._array.__abs__()
         return self.__class__._new(res, device=self.device)
 
-    def __add__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __add__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __add__.
         """
@@ -530,7 +558,7 @@ def __add__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__add__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __and__(self: Array, other: Union[int, bool, Array], /) -> Array:
+    def __and__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __and__.
         """
@@ -542,9 +570,7 @@ def __and__(self: Array, other: Union[int, bool, Array], /) -> Array:
         res = self._array.__and__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __array_namespace__(
-        self: Array, /, *, api_version: Optional[str] = None
-    ) -> types.ModuleType:
+    def __array_namespace__(self, /, *, api_version: str | None = None) -> ModuleType:
         """
         Return the array_api_strict namespace corresponding to api_version.
 
@@ -563,7 +589,7 @@ def __array_namespace__(
         import array_api_strict
         return array_api_strict
 
-    def __bool__(self: Array, /) -> bool:
+    def __bool__(self) -> bool:
         """
         Performs the operation __bool__.
         """
@@ -573,7 +599,7 @@ def __bool__(self: Array, /) -> bool:
         res = self._array.__bool__()
         return res
 
-    def __complex__(self: Array, /) -> complex:
+    def __complex__(self) -> complex:
         """
         Performs the operation __complex__.
         """
@@ -584,52 +610,52 @@ def __complex__(self: Array, /) -> complex:
         return res
 
     def __dlpack__(
-        self: Array,
+        self,
         /,
         *,
-        stream: Optional[Union[int, Any]] = None,
-        max_version: Optional[tuple[int, int]] = _default,
-        dl_device: Optional[tuple[IntEnum, int]] = _default,
-        copy: Optional[bool] = _default,
+        stream: Any = None,
+        max_version: tuple[int, int] | None | Undef = _undef,
+        dl_device: tuple[IntEnum, int] | None | Undef = _undef,
+        copy: bool | None | Undef = _undef,
     ) -> PyCapsule:
         """
         Performs the operation __dlpack__.
         """
         if get_array_api_strict_flags()['api_version'] < '2023.12':
-            if max_version is not _default:
+            if max_version is not _undef:
                 raise ValueError("The max_version argument to __dlpack__ requires at least version 2023.12 of the array API")
-            if dl_device is not _default:
+            if dl_device is not _undef:
                 raise ValueError("The device argument to __dlpack__ requires at least version 2023.12 of the array API")
-            if copy is not _default:
+            if copy is not _undef:
                 raise ValueError("The copy argument to __dlpack__ requires at least version 2023.12 of the array API")
 
         if np.lib.NumpyVersion(np.__version__) < '2.1.0':
-            if max_version not in [_default, None]:
+            if max_version not in [_undef, None]:
                 raise NotImplementedError("The max_version argument to __dlpack__ is not yet implemented")
-            if dl_device not in [_default, None]:
+            if dl_device not in [_undef, None]:
                 raise NotImplementedError("The device argument to __dlpack__ is not yet implemented")
-            if copy not in [_default, None]:
+            if copy not in [_undef, None]:
                 raise NotImplementedError("The copy argument to __dlpack__ is not yet implemented")
 
             return self._array.__dlpack__(stream=stream)
         else:
             kwargs = {'stream': stream}
-            if max_version is not _default:
+            if max_version is not _undef:
                 kwargs['max_version'] = max_version
-            if dl_device is not _default:
+            if dl_device is not _undef:
                 kwargs['dl_device'] = dl_device
-            if copy is not _default:
+            if copy is not _undef:
                 kwargs['copy'] = copy
             return self._array.__dlpack__(**kwargs)
 
-    def __dlpack_device__(self: Array, /) -> Tuple[IntEnum, int]:
+    def __dlpack_device__(self) -> tuple[IntEnum, int]:
         """
         Performs the operation __dlpack_device__.
         """
         # Note: device support is required for this
         return self._array.__dlpack_device__()
 
-    def __eq__(self: Array, other: Union[int, float, bool, Array], /) -> Array:
+    def __eq__(self, other: Array | bool | int | float | complex, /) -> Array:  # type: ignore[override]
         """
         Performs the operation __eq__.
         """
@@ -643,7 +669,7 @@ def __eq__(self: Array, other: Union[int, float, bool, Array], /) -> Array:
         res = self._array.__eq__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __float__(self: Array, /) -> float:
+    def __float__(self) -> float:
         """
         Performs the operation __float__.
         """
@@ -655,7 +681,7 @@ def __float__(self: Array, /) -> float:
         res = self._array.__float__()
         return res
 
-    def __floordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __floordiv__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __floordiv__.
         """
@@ -667,7 +693,7 @@ def __floordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__floordiv__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ge__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __ge__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __ge__.
         """
@@ -680,14 +706,15 @@ def __ge__(self: Array, other: Union[int, float, Array], /) -> Array:
         return self.__class__._new(res, device=self.device)
 
     def __getitem__(
-        self: Array,
-        key: Union[
-            int,
-            slice,
-            ellipsis,  # noqa: F821
-            Tuple[Union[int, slice, ellipsis, None], ...],  # noqa: F821
-            Array,
-        ],
+        self,
+        key: (
+            int
+            | slice
+            | EllipsisType
+            | Array
+            | None
+            | tuple[int | slice | EllipsisType | Array | None, ...]
+        ),
         /,
     ) -> Array:
         """
@@ -696,14 +723,13 @@ def __getitem__(
         # XXX Does key have to be on the same device? Is there an exception for CPU_DEVICE?
         # Note: Only indices required by the spec are allowed. See the
         # docstring of _validate_index
-        self._validate_index(key)
-        if isinstance(key, Array):
-            # Indexing self._array with array_api_strict arrays can be erroneous
-            key = key._array
-        res = self._array.__getitem__(key)
+        self._validate_index(key, op="getitem")
+        # Indexing self._array with array_api_strict arrays can be erroneous
+        np_key = key._array if isinstance(key, Array) else key
+        res = self._array.__getitem__(np_key)
         return self._new(res, device=self.device)
 
-    def __gt__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __gt__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __gt__.
         """
@@ -715,7 +741,7 @@ def __gt__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__gt__(other._array)
         return self.__class__._new(res, device=other.device)
 
-    def __int__(self: Array, /) -> int:
+    def __int__(self) -> int:
         """
         Performs the operation __int__.
         """
@@ -727,14 +753,14 @@ def __int__(self: Array, /) -> int:
         res = self._array.__int__()
         return res
 
-    def __index__(self: Array, /) -> int:
+    def __index__(self) -> int:
         """
         Performs the operation __index__.
         """
         res = self._array.__index__()
         return res
 
-    def __invert__(self: Array, /) -> Array:
+    def __invert__(self) -> Array:
         """
         Performs the operation __invert__.
         """
@@ -743,7 +769,7 @@ def __invert__(self: Array, /) -> Array:
         res = self._array.__invert__()
         return self.__class__._new(res, device=self.device)
 
-    def __iter__(self: Array, /):
+    def __iter__(self) -> Iterator[Array]:
         """
         Performs the operation __iter__.
         """
@@ -758,7 +784,7 @@ def __iter__(self: Array, /):
         # implemented, which implies iteration on 1-D arrays.
         return (Array._new(i, device=self.device) for i in self._array)
 
-    def __le__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __le__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __le__.
         """
@@ -770,7 +796,7 @@ def __le__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__le__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __lshift__(self: Array, other: Union[int, Array], /) -> Array:
+    def __lshift__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __lshift__.
         """
@@ -782,7 +808,7 @@ def __lshift__(self: Array, other: Union[int, Array], /) -> Array:
         res = self._array.__lshift__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __lt__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __lt__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __lt__.
         """
@@ -794,7 +820,7 @@ def __lt__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__lt__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __matmul__(self: Array, other: Array, /) -> Array:
+    def __matmul__(self, other: Array, /) -> Array:
         """
         Performs the operation __matmul__.
         """
@@ -807,7 +833,7 @@ def __matmul__(self: Array, other: Array, /) -> Array:
         res = self._array.__matmul__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __mod__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __mod__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __mod__.
         """
@@ -819,7 +845,7 @@ def __mod__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__mod__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __mul__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __mul__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __mul__.
         """
@@ -831,7 +857,7 @@ def __mul__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__mul__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ne__(self: Array, other: Union[int, float, bool, Array], /) -> Array:
+    def __ne__(self, other: Array | bool | int | float | complex, /) -> Array:  # type: ignore[override]
         """
         Performs the operation __ne__.
         """
@@ -843,7 +869,7 @@ def __ne__(self: Array, other: Union[int, float, bool, Array], /) -> Array:
         res = self._array.__ne__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __neg__(self: Array, /) -> Array:
+    def __neg__(self) -> Array:
         """
         Performs the operation __neg__.
         """
@@ -852,7 +878,7 @@ def __neg__(self: Array, /) -> Array:
         res = self._array.__neg__()
         return self.__class__._new(res, device=self.device)
 
-    def __or__(self: Array, other: Union[int, bool, Array], /) -> Array:
+    def __or__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __or__.
         """
@@ -864,7 +890,7 @@ def __or__(self: Array, other: Union[int, bool, Array], /) -> Array:
         res = self._array.__or__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __pos__(self: Array, /) -> Array:
+    def __pos__(self) -> Array:
         """
         Performs the operation __pos__.
         """
@@ -873,11 +899,11 @@ def __pos__(self: Array, /) -> Array:
         res = self._array.__pos__()
         return self.__class__._new(res, device=self.device)
 
-    def __pow__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __pow__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __pow__.
         """
-        from ._elementwise_functions import pow
+        from ._elementwise_functions import pow  # type: ignore[attr-defined]
 
         self._check_device(other)
         other = self._check_allowed_dtypes(other, "numeric", "__pow__")
@@ -887,7 +913,7 @@ def __pow__(self: Array, other: Union[int, float, Array], /) -> Array:
         # arrays, so we use pow() here instead.
         return pow(self, other)
 
-    def __rshift__(self: Array, other: Union[int, Array], /) -> Array:
+    def __rshift__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __rshift__.
         """
@@ -901,10 +927,16 @@ def __rshift__(self: Array, other: Union[int, Array], /) -> Array:
 
     def __setitem__(
         self,
-        key: Union[
-            int, slice, ellipsis, Tuple[Union[int, slice, ellipsis], ...], Array  # noqa: F821
-        ],
-        value: Union[int, float, bool, Array],
+        # Almost same as __getitem__ key but doesn't accept None
+        # or integer arrays
+        key: (
+            int
+            | slice
+            | EllipsisType
+            | Array
+            | tuple[int | slice | EllipsisType, ...]
+        ),
+        value: Array | bool | int | float | complex,
         /,
     ) -> None:
         """
@@ -913,12 +945,11 @@ def __setitem__(
         # Note: Only indices required by the spec are allowed. See the
         # docstring of _validate_index
         self._validate_index(key, op="setitem")
-        if isinstance(key, Array):
-            # Indexing self._array with array_api_strict arrays can be erroneous
-            key = key._array
-        self._array.__setitem__(key, asarray(value)._array)
+        # Indexing self._array with array_api_strict arrays can be erroneous
+        np_key = key._array if isinstance(key, Array) else key
+        self._array.__setitem__(np_key, asarray(value)._array)
 
-    def __sub__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __sub__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __sub__.
         """
@@ -932,7 +963,7 @@ def __sub__(self: Array, other: Union[int, float, Array], /) -> Array:
 
     # PEP 484 requires int to be a subtype of float, but __truediv__ should
     # not accept int.
-    def __truediv__(self: Array, other: Union[float, Array], /) -> Array:
+    def __truediv__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __truediv__.
         """
@@ -944,7 +975,7 @@ def __truediv__(self: Array, other: Union[float, Array], /) -> Array:
         res = self._array.__truediv__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __xor__(self: Array, other: Union[int, bool, Array], /) -> Array:
+    def __xor__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __xor__.
         """
@@ -956,7 +987,7 @@ def __xor__(self: Array, other: Union[int, bool, Array], /) -> Array:
         res = self._array.__xor__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __iadd__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __iadd__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __iadd__.
         """
@@ -967,7 +998,7 @@ def __iadd__(self: Array, other: Union[int, float, Array], /) -> Array:
         self._array.__iadd__(other._array)
         return self
 
-    def __radd__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __radd__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __radd__.
         """
@@ -979,7 +1010,7 @@ def __radd__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__radd__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __iand__(self: Array, other: Union[int, bool, Array], /) -> Array:
+    def __iand__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __iand__.
         """
@@ -990,7 +1021,7 @@ def __iand__(self: Array, other: Union[int, bool, Array], /) -> Array:
         self._array.__iand__(other._array)
         return self
 
-    def __rand__(self: Array, other: Union[int, bool, Array], /) -> Array:
+    def __rand__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __rand__.
         """
@@ -1002,7 +1033,7 @@ def __rand__(self: Array, other: Union[int, bool, Array], /) -> Array:
         res = self._array.__rand__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ifloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __ifloordiv__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __ifloordiv__.
         """
@@ -1013,7 +1044,7 @@ def __ifloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
         self._array.__ifloordiv__(other._array)
         return self
 
-    def __rfloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __rfloordiv__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __rfloordiv__.
         """
@@ -1025,7 +1056,7 @@ def __rfloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__rfloordiv__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ilshift__(self: Array, other: Union[int, Array], /) -> Array:
+    def __ilshift__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __ilshift__.
         """
@@ -1036,7 +1067,7 @@ def __ilshift__(self: Array, other: Union[int, Array], /) -> Array:
         self._array.__ilshift__(other._array)
         return self
 
-    def __rlshift__(self: Array, other: Union[int, Array], /) -> Array:
+    def __rlshift__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __rlshift__.
         """
@@ -1048,7 +1079,7 @@ def __rlshift__(self: Array, other: Union[int, Array], /) -> Array:
         res = self._array.__rlshift__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __imatmul__(self: Array, other: Array, /) -> Array:
+    def __imatmul__(self, other: Array, /) -> Array:
         """
         Performs the operation __imatmul__.
         """
@@ -1061,7 +1092,7 @@ def __imatmul__(self: Array, other: Array, /) -> Array:
         res = self._array.__imatmul__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __rmatmul__(self: Array, other: Array, /) -> Array:
+    def __rmatmul__(self, other: Array, /) -> Array:
         """
         Performs the operation __rmatmul__.
         """
@@ -1074,7 +1105,7 @@ def __rmatmul__(self: Array, other: Array, /) -> Array:
         res = self._array.__rmatmul__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __imod__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __imod__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __imod__.
         """
@@ -1084,7 +1115,7 @@ def __imod__(self: Array, other: Union[int, float, Array], /) -> Array:
         self._array.__imod__(other._array)
         return self
 
-    def __rmod__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __rmod__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __rmod__.
         """
@@ -1096,7 +1127,7 @@ def __rmod__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__rmod__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __imul__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __imul__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __imul__.
         """
@@ -1106,7 +1137,7 @@ def __imul__(self: Array, other: Union[int, float, Array], /) -> Array:
         self._array.__imul__(other._array)
         return self
 
-    def __rmul__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __rmul__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __rmul__.
         """
@@ -1118,7 +1149,7 @@ def __rmul__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__rmul__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ior__(self: Array, other: Union[int, bool, Array], /) -> Array:
+    def __ior__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __ior__.
         """
@@ -1128,7 +1159,7 @@ def __ior__(self: Array, other: Union[int, bool, Array], /) -> Array:
         self._array.__ior__(other._array)
         return self
 
-    def __ror__(self: Array, other: Union[int, bool, Array], /) -> Array:
+    def __ror__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __ror__.
         """
@@ -1140,7 +1171,7 @@ def __ror__(self: Array, other: Union[int, bool, Array], /) -> Array:
         res = self._array.__ror__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ipow__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __ipow__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __ipow__.
         """
@@ -1150,11 +1181,11 @@ def __ipow__(self: Array, other: Union[int, float, Array], /) -> Array:
         self._array.__ipow__(other._array)
         return self
 
-    def __rpow__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __rpow__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __rpow__.
         """
-        from ._elementwise_functions import pow
+        from ._elementwise_functions import pow  # type: ignore[attr-defined]
 
         other = self._check_allowed_dtypes(other, "numeric", "__rpow__")
         if other is NotImplemented:
@@ -1163,7 +1194,7 @@ def __rpow__(self: Array, other: Union[int, float, Array], /) -> Array:
         # for 0-d arrays, so we use pow() here instead.
         return pow(other, self)
 
-    def __irshift__(self: Array, other: Union[int, Array], /) -> Array:
+    def __irshift__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __irshift__.
         """
@@ -1173,7 +1204,7 @@ def __irshift__(self: Array, other: Union[int, Array], /) -> Array:
         self._array.__irshift__(other._array)
         return self
 
-    def __rrshift__(self: Array, other: Union[int, Array], /) -> Array:
+    def __rrshift__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __rrshift__.
         """
@@ -1185,7 +1216,7 @@ def __rrshift__(self: Array, other: Union[int, Array], /) -> Array:
         res = self._array.__rrshift__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __isub__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __isub__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __isub__.
         """
@@ -1195,7 +1226,7 @@ def __isub__(self: Array, other: Union[int, float, Array], /) -> Array:
         self._array.__isub__(other._array)
         return self
 
-    def __rsub__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __rsub__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __rsub__.
         """
@@ -1207,7 +1238,7 @@ def __rsub__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__rsub__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __itruediv__(self: Array, other: Union[float, Array], /) -> Array:
+    def __itruediv__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __itruediv__.
         """
@@ -1217,7 +1248,7 @@ def __itruediv__(self: Array, other: Union[float, Array], /) -> Array:
         self._array.__itruediv__(other._array)
         return self
 
-    def __rtruediv__(self: Array, other: Union[float, Array], /) -> Array:
+    def __rtruediv__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __rtruediv__.
         """
@@ -1229,7 +1260,7 @@ def __rtruediv__(self: Array, other: Union[float, Array], /) -> Array:
         res = self._array.__rtruediv__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ixor__(self: Array, other: Union[int, bool, Array], /) -> Array:
+    def __ixor__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __ixor__.
         """
@@ -1239,7 +1270,7 @@ def __ixor__(self: Array, other: Union[int, bool, Array], /) -> Array:
         self._array.__ixor__(other._array)
         return self
 
-    def __rxor__(self: Array, other: Union[int, bool, Array], /) -> Array:
+    def __rxor__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __rxor__.
         """
@@ -1251,7 +1282,7 @@ def __rxor__(self: Array, other: Union[int, bool, Array], /) -> Array:
         res = self._array.__rxor__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def to_device(self: Array, device: Device, /, stream: None = None) -> Array:
+    def to_device(self, device: Device, /, stream: None = None) -> Array:
         if stream is not None:
             raise ValueError("The stream argument to to_device() is not supported")
         if device == self._device:
@@ -1262,7 +1293,7 @@ def to_device(self: Array, device: Device, /, stream: None = None) -> Array:
         raise ValueError(f"Unsupported device {device!r}")
 
     @property
-    def dtype(self) -> Dtype:
+    def dtype(self) -> DType:
         """
         Array API compatible wrapper for :py:meth:`np.ndarray.dtype <numpy.ndarray.dtype>`.
 
@@ -1290,7 +1321,7 @@ def ndim(self) -> int:
         return self._array.ndim
 
     @property
-    def shape(self) -> Tuple[int, ...]:
+    def shape(self) -> tuple[int, ...]:
         """
         Array API compatible wrapper for :py:meth:`np.ndarray.shape <numpy.ndarray.shape>`.
 
diff --git a/array_api_strict/_constants.py b/array_api_strict/_constants.py
index 15ab81d..d78354b 100644
--- a/array_api_strict/_constants.py
+++ b/array_api_strict/_constants.py
@@ -4,4 +4,4 @@
 inf = np.inf
 nan = np.nan
 pi = np.pi
-newaxis = np.newaxis
+newaxis: None = np.newaxis
diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py
index 460dba9..3b80b8a 100644
--- a/array_api_strict/_creation_functions.py
+++ b/array_api_strict/_creation_functions.py
@@ -1,23 +1,33 @@
 from __future__ import annotations
 
+from collections.abc import Generator
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+from enum import Enum
+from typing import TYPE_CHECKING, Literal
 
-if TYPE_CHECKING:
-    from ._typing import (
-        Array,
-        Device,
-        Dtype,
-        NestedSequence,
-        SupportsBufferProtocol,
-    )
-from ._dtypes import _DType, _all_dtypes
+import numpy as np
+
+from ._dtypes import DType, _all_dtypes, _np_dtype
 from ._flags import get_array_api_strict_flags
+from ._typing import NestedSequence, SupportsBufferProtocol, SupportsDLPack
+
+if TYPE_CHECKING:
+    # TODO import from typing (requires Python >=3.13)
+    from typing_extensions import TypeIs
+
+    # Circular import
+    from ._array_object import Array, Device
+
+
+class Undef(Enum):
+    UNDEF = 0
+
+
+_undef = Undef.UNDEF
 
-import numpy as np
 
 @contextmanager
-def allow_array():
+def allow_array() -> Generator[None]:
     """
     Temporarily enable Array.__array__. This is needed for np.array to parse
     list of lists of Array objects.
@@ -30,22 +40,25 @@ def allow_array():
     finally:
         _array_object._allow_array = original_value
 
-def _check_valid_dtype(dtype):
+
+def _check_valid_dtype(dtype: DType | None) -> None:
     # Note: Only spelling dtypes as the dtype objects is supported.
     if dtype not in (None,) + _all_dtypes:
         raise ValueError(f"dtype must be one of the supported dtypes, got {dtype!r}")
 
-def _supports_buffer_protocol(obj):
+
+def _supports_buffer_protocol(obj: object) -> TypeIs[SupportsBufferProtocol]:
     try:
-        memoryview(obj)
+        memoryview(obj)  # type: ignore[arg-type]
     except TypeError:
         return False
     return True
 
-def _check_device(device):
+
+def _check_device(device: Device | None) -> None:
     # _array_object imports in this file are inside the functions to avoid
     # circular imports
-    from ._array_object import Device, ALL_DEVICES
+    from ._array_object import ALL_DEVICES, Device
 
     if device is not None and not isinstance(device, Device):
         raise ValueError(f"Unsupported device {device!r}")
@@ -53,20 +66,20 @@ def _check_device(device):
     if device is not None and device not in ALL_DEVICES:
         raise ValueError(f"Unsupported device {device!r}")
 
+
 def asarray(
-    obj: Union[
-        Array,
-        bool,
-        int,
-        float,
-        NestedSequence[bool | int | float],
-        SupportsBufferProtocol,
-    ],
+    obj: Array
+    | bool
+    | int
+    | float
+    | complex
+    | NestedSequence[bool | int | float | complex]
+    | SupportsBufferProtocol,
     /,
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
-    copy: Optional[bool] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
+    copy: bool | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.asarray <numpy.asarray>`.
@@ -118,13 +131,13 @@ def asarray(
 
 
 def arange(
-    start: Union[int, float],
+    start: int | float,
     /,
-    stop: Optional[Union[int, float]] = None,
-    step: Union[int, float] = 1,
+    stop: int | float | None = None,
+    step: int | float = 1,
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.arange <numpy.arange>`.
@@ -136,16 +149,17 @@ def arange(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.arange(start, stop=stop, step=step, dtype=dtype), device=device)
+    return Array._new(
+        np.arange(start, stop, step, dtype=_np_dtype(dtype)),
+        device=device,
+    )
 
 
 def empty(
-    shape: Union[int, Tuple[int, ...]],
+    shape: int | tuple[int, ...],
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.empty <numpy.empty>`.
@@ -157,13 +171,11 @@ def empty(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.empty(shape, dtype=dtype), device=device)
+    return Array._new(np.empty(shape, dtype=_np_dtype(dtype)), device=device)
 
 
 def empty_like(
-    x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None
+    x: Array, /, *, dtype: DType | None = None, device: Device | None = None
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.empty_like <numpy.empty_like>`.
@@ -177,19 +189,17 @@ def empty_like(
     if device is None:
         device = x.device
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.empty_like(x._array, dtype=dtype), device=device)
+    return Array._new(np.empty_like(x._array, dtype=_np_dtype(dtype)), device=device)
 
 
 def eye(
     n_rows: int,
-    n_cols: Optional[int] = None,
+    n_cols: int | None = None,
     /,
     *,
     k: int = 0,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.eye <numpy.eye>`.
@@ -201,45 +211,43 @@ def eye(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.eye(n_rows, M=n_cols, k=k, dtype=dtype), device=device)
-
+    return Array._new(
+        np.eye(n_rows, M=n_cols, k=k, dtype=_np_dtype(dtype)), device=device
+    )
 
-_default = object()
 
 def from_dlpack(
-    x: object,
+    x: SupportsDLPack,
     /,
     *,
-    device: Optional[Device] = _default,
-    copy: Optional[bool] = _default,
+    device: Device | Undef | None = _undef,
+    copy: bool | Undef | None = _undef,
 ) -> Array:
     from ._array_object import Array
 
     if get_array_api_strict_flags()['api_version'] < '2023.12':
-        if device is not _default:
+        if device is not _undef:
             raise ValueError("The device argument to from_dlpack requires at least version 2023.12 of the array API")
-        if copy is not _default:
+        if copy is not _undef:
             raise ValueError("The copy argument to from_dlpack requires at least version 2023.12 of the array API")
 
     # Going to wait for upstream numpy support
-    if device is not _default:
+    if device is not _undef:
         _check_device(device)
     else:
         device = None
-    if copy not in [_default, None]:
+    if copy not in [_undef, None]:
         raise NotImplementedError("The copy argument to from_dlpack is not yet implemented")
 
     return Array._new(np.from_dlpack(x), device=device)
 
 
 def full(
-    shape: Union[int, Tuple[int, ...]],
-    fill_value: Union[int, float],
+    shape: int | tuple[int, ...],
+    fill_value: bool | int | float | complex,
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.full <numpy.full>`.
@@ -253,10 +261,8 @@ def full(
 
     if isinstance(fill_value, Array) and fill_value.ndim == 0:
         fill_value = fill_value._array
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    res = np.full(shape, fill_value, dtype=dtype)
-    if _DType(res.dtype) not in _all_dtypes:
+    res = np.full(shape, fill_value, dtype=_np_dtype(dtype))
+    if DType(res.dtype) not in _all_dtypes:
         # This will happen if the fill value is not something that NumPy
         # coerces to one of the acceptable dtypes.
         raise TypeError("Invalid input to full")
@@ -266,10 +272,10 @@ def full(
 def full_like(
     x: Array,
     /,
-    fill_value: Union[int, float],
+    fill_value: bool | int | float | complex,
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.full_like <numpy.full_like>`.
@@ -283,10 +289,8 @@ def full_like(
     if device is None:
         device = x.device
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    res = np.full_like(x._array, fill_value, dtype=dtype)
-    if _DType(res.dtype) not in _all_dtypes:
+    res = np.full_like(x._array, fill_value, dtype=_np_dtype(dtype))
+    if DType(res.dtype) not in _all_dtypes:
         # This will happen if the fill value is not something that NumPy
         # coerces to one of the acceptable dtypes.
         raise TypeError("Invalid input to full_like")
@@ -294,13 +298,13 @@ def full_like(
 
 
 def linspace(
-    start: Union[int, float],
-    stop: Union[int, float],
+    start: int | float | complex,
+    stop: int | float | complex,
     /,
     num: int,
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
     endpoint: bool = True,
 ) -> Array:
     """
@@ -313,12 +317,13 @@ def linspace(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.linspace(start, stop, num, dtype=dtype, endpoint=endpoint), device=device)
+    return Array._new(
+        np.linspace(start, stop, num, dtype=_np_dtype(dtype), endpoint=endpoint),
+        device=device,
+    )
 
 
-def meshgrid(*arrays: Array, indexing: str = "xy") -> List[Array]:
+def meshgrid(*arrays: Array, indexing: Literal["xy", "ij"] = "xy") -> list[Array]:
     """
     Array API compatible wrapper for :py:func:`np.meshgrid <numpy.meshgrid>`.
 
@@ -348,10 +353,10 @@ def meshgrid(*arrays: Array, indexing: str = "xy") -> List[Array]:
 
 
 def ones(
-    shape: Union[int, Tuple[int, ...]],
+    shape: int | tuple[int, ...],
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.ones <numpy.ones>`.
@@ -363,13 +368,11 @@ def ones(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.ones(shape, dtype=dtype), device=device)
+    return Array._new(np.ones(shape, dtype=_np_dtype(dtype)), device=device)
 
 
 def ones_like(
-    x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None
+    x: Array, /, *, dtype: DType | None = None, device: Device | None = None
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.ones_like <numpy.ones_like>`.
@@ -383,9 +386,7 @@ def ones_like(
     if device is None:
         device = x.device
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.ones_like(x._array, dtype=dtype), device=device)
+    return Array._new(np.ones_like(x._array, dtype=_np_dtype(dtype)), device=device)
 
 
 def tril(x: Array, /, *, k: int = 0) -> Array:
@@ -417,10 +418,10 @@ def triu(x: Array, /, *, k: int = 0) -> Array:
 
 
 def zeros(
-    shape: Union[int, Tuple[int, ...]],
+    shape: int | tuple[int, ...],
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.zeros <numpy.zeros>`.
@@ -432,13 +433,11 @@ def zeros(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.zeros(shape, dtype=dtype), device=device)
+    return Array._new(np.zeros(shape, dtype=_np_dtype(dtype)), device=device)
 
 
 def zeros_like(
-    x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None
+    x: Array, /, *, dtype: DType | None = None, device: Device | None = None
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.zeros_like <numpy.zeros_like>`.
@@ -452,6 +451,4 @@ def zeros_like(
     if device is None:
         device = x.device
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.zeros_like(x._array, dtype=dtype), device=device)
+    return Array._new(np.zeros_like(x._array, dtype=_np_dtype(dtype)), device=device)
diff --git a/array_api_strict/_data_type_functions.py b/array_api_strict/_data_type_functions.py
index 1643043..7dc918d 100644
--- a/array_api_strict/_data_type_functions.py
+++ b/array_api_strict/_data_type_functions.py
@@ -1,38 +1,37 @@
 from __future__ import annotations
 
-from ._array_object import Array
-from ._creation_functions import _check_device
+from dataclasses import dataclass
+
+import numpy as np
+
+from ._array_object import Array, Device
+from ._creation_functions import Undef, _check_device, _undef
 from ._dtypes import (
-    _DType,
+    DType,
     _all_dtypes,
     _boolean_dtypes,
-    _signed_integer_dtypes,
-    _unsigned_integer_dtypes,
-    _integer_dtypes,
-    _real_floating_dtypes,
     _complex_floating_dtypes,
+    _integer_dtypes,
     _numeric_dtypes,
+    _real_floating_dtypes,
     _result_type,
+    _signed_integer_dtypes,
+    _unsigned_integer_dtypes,
 )
 from ._flags import get_array_api_strict_flags
 
-from dataclasses import dataclass
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from typing import List, Tuple, Union, Optional
-    from ._typing import Dtype, Device
-
-import numpy as np
-
-# Use to emulate the asarray(device) argument not existing in 2022.12
-_default = object()
 
 # Note: astype is a function, not an array method as in NumPy.
 def astype(
-    x: Array, dtype: Dtype, /, *, copy: bool = True, device: Optional[Device] = _default
+    x: Array,
+    dtype: DType,
+    /,
+    *,
+    copy: bool = True,
+    # _default is used to emulate the device argument not existing in 2022.12
+    device: Device | Undef | None = _undef,
 ) -> Array:
-    if device is not _default:
+    if device is not _undef:
         if get_array_api_strict_flags()['api_version'] >= '2023.12':
             _check_device(device)
         else:
@@ -52,7 +51,7 @@ def astype(
     return Array._new(x._array.astype(dtype=dtype._np_dtype, copy=copy), device=device)
 
 
-def broadcast_arrays(*arrays: Array) -> List[Array]:
+def broadcast_arrays(*arrays: Array) -> list[Array]:
     """
     Array API compatible wrapper for :py:func:`np.broadcast_arrays <numpy.broadcast_arrays>`.
 
@@ -65,7 +64,7 @@ def broadcast_arrays(*arrays: Array) -> List[Array]:
     ]
 
 
-def broadcast_to(x: Array, /, shape: Tuple[int, ...]) -> Array:
+def broadcast_to(x: Array, /, shape: tuple[int, ...]) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.broadcast_to <numpy.broadcast_to>`.
 
@@ -76,7 +75,7 @@ def broadcast_to(x: Array, /, shape: Tuple[int, ...]) -> Array:
     return Array._new(np.broadcast_to(x._array, shape), device=x.device)
 
 
-def can_cast(from_: Union[Dtype, Array], to: Dtype, /) -> bool:
+def can_cast(from_: DType | Array, to: DType, /) -> bool:
     """
     Array API compatible wrapper for :py:func:`np.can_cast <numpy.can_cast>`.
 
@@ -112,7 +111,7 @@ class finfo_object:
     max: float
     min: float
     smallest_normal: float
-    dtype: Dtype
+    dtype: DType
 
 
 @dataclass
@@ -120,18 +119,17 @@ class iinfo_object:
     bits: int
     max: int
     min: int
-    dtype: Dtype
+    dtype: DType
 
 
-def finfo(type: Union[Dtype, Array], /) -> finfo_object:
+def finfo(type: DType | Array, /) -> finfo_object:
     """
     Array API compatible wrapper for :py:func:`np.finfo <numpy.finfo>`.
 
     See its docstring for more information.
     """
-    if isinstance(type, _DType):
-        type = type._np_dtype
-    fi = np.finfo(type)
+    np_type = type._array if isinstance(type, Array) else type._np_dtype
+    fi = np.finfo(np_type)
     # Note: The types of the float data here are float, whereas in NumPy they
     # are scalars of the corresponding float dtype.
     return finfo_object(
@@ -140,35 +138,33 @@ def finfo(type: Union[Dtype, Array], /) -> finfo_object:
         float(fi.max),
         float(fi.min),
         float(fi.smallest_normal),
-        fi.dtype,
+        DType(fi.dtype),
     )
 
 
-def iinfo(type: Union[Dtype, Array], /) -> iinfo_object:
+def iinfo(type: DType | Array, /) -> iinfo_object:
     """
     Array API compatible wrapper for :py:func:`np.iinfo <numpy.iinfo>`.
 
     See its docstring for more information.
     """
-    if isinstance(type, _DType):
-        type = type._np_dtype
-    ii = np.iinfo(type)
-    return iinfo_object(ii.bits, ii.max, ii.min, ii.dtype)
+    np_type = type._array if isinstance(type, Array) else type._np_dtype
+    ii = np.iinfo(np_type)
+    return iinfo_object(ii.bits, ii.max, ii.min, DType(ii.dtype))
 
 
 # Note: isdtype is a new function from the 2022.12 array API specification.
-def isdtype(
-    dtype: Dtype, kind: Union[Dtype, str, Tuple[Union[Dtype, str], ...]]
-) -> bool:
+def isdtype(dtype: DType, kind: DType | str | tuple[DType | str, ...]) -> bool:
     """
-    Returns a boolean indicating whether a provided dtype is of a specified data type ``kind``.
+    Returns a boolean indicating whether a provided dtype is of a specified
+    data type ``kind``.
 
     See
     https://data-apis.org/array-api/latest/API_specification/generated/array_api.isdtype.html
     for more details
     """
-    if not isinstance(dtype, _DType):
-       raise TypeError(f"'dtype' must be a dtype, not a {type(dtype)!r}")
+    if not isinstance(dtype, DType):
+        raise TypeError(f"'dtype' must be a dtype, not a {type(dtype)!r}")
 
     if isinstance(kind, tuple):
         # Disallow nested tuples
@@ -197,7 +193,10 @@ def isdtype(
     else:
         raise TypeError(f"'kind' must be a dtype, str, or tuple of dtypes and strs, not {type(kind).__name__}")
 
-def result_type(*arrays_and_dtypes: Union[Array, Dtype, int, float, complex, bool]) -> Dtype:
+
+def result_type(
+    *arrays_and_dtypes: DType | Array | bool | int | float | complex,
+) -> DType:
     """
     Array API compatible wrapper for :py:func:`np.result_type <numpy.result_type>`.
 
@@ -219,15 +218,15 @@ def result_type(*arrays_and_dtypes: Union[Array, Dtype, int, float, complex, boo
         A.append(a)
 
     # remove python scalars
-    A = [a for a in A if not isinstance(a, (bool, int, float, complex))]
+    B = [a for a in A if not isinstance(a, (bool, int, float, complex))]
 
-    if len(A) == 0:
+    if len(B) == 0:
         raise ValueError("at least one array or dtype is required")
-    elif len(A) == 1:
-        result = A[0]
+    elif len(B) == 1:
+        result = B[0]
     else:
-        t = A[0]
-        for t2 in A[1:]:
+        t = B[0]
+        for t2 in B[1:]:
             t = _result_type(t, t2)
         result = t
 
diff --git a/array_api_strict/_dtypes.py b/array_api_strict/_dtypes.py
index 66304dd..513650b 100644
--- a/array_api_strict/_dtypes.py
+++ b/array_api_strict/_dtypes.py
@@ -1,19 +1,27 @@
+from __future__ import annotations
+
+import builtins
 import warnings
+from typing import Any, Final
 
 import numpy as np
+import numpy.typing as npt
 
 # Note: we wrap the NumPy dtype objects in a bare class, so that none of the
 # additional methods and behaviors of NumPy dtype objects are exposed.
 
-class _DType:
-    def __init__(self, np_dtype):
-        np_dtype = np.dtype(np_dtype)
-        self._np_dtype = np_dtype
 
-    def __repr__(self):
+class DType:
+    _np_dtype: Final[np.dtype[Any]]
+    __slots__ = ("_np_dtype", "__weakref__")
+
+    def __init__(self, np_dtype: npt.DTypeLike):
+        self._np_dtype = np.dtype(np_dtype)
+
+    def __repr__(self) -> str:
         return f"array_api_strict.{self._np_dtype.name}"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> builtins.bool:
         # See https://github.com/numpy/numpy/pull/25370/files#r1423259515.
         # Avoid the user error of array_api_strict.float32 == numpy.float32,
         # which gives False. Making == error is probably too egregious, so
@@ -26,12 +34,13 @@ def __eq__(self, other):
 a NumPy native dtype object, but you probably don't want to do this. \
 array_api_strict dtype objects compare unequal to their NumPy equivalents. \
 Such cross-library comparison is not supported by the standard.""",
-            stacklevel=2)
-        if not isinstance(other, _DType):
+                stacklevel=2,
+            )
+        if not isinstance(other, DType):
             return NotImplemented
         return self._np_dtype == other._np_dtype
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         # Note: this is not strictly required
         # (https://github.com/data-apis/array-api/issues/582), but makes the
         # dtype objects much easier to work with here and elsewhere if they
@@ -39,20 +48,24 @@ def __hash__(self):
         return hash(self._np_dtype)
 
 
-int8 = _DType("int8")
-int16 = _DType("int16")
-int32 = _DType("int32")
-int64 = _DType("int64")
-uint8 = _DType("uint8")
-uint16 = _DType("uint16")
-uint32 = _DType("uint32")
-uint64 = _DType("uint64")
-float32 = _DType("float32")
-float64 = _DType("float64")
-complex64 = _DType("complex64")
-complex128 = _DType("complex128")
+def _np_dtype(dtype: DType | None) -> np.dtype[Any] | None:
+    return dtype._np_dtype if dtype is not None else None
+
+
+int8 = DType("int8")
+int16 = DType("int16")
+int32 = DType("int32")
+int64 = DType("int64")
+uint8 = DType("uint8")
+uint16 = DType("uint16")
+uint32 = DType("uint32")
+uint64 = DType("uint64")
+float32 = DType("float32")
+float64 = DType("float64")
+complex64 = DType("complex64")
+complex128 = DType("complex128")
 # Note: This name is changed
-bool = _DType("bool")
+bool = DType("bool")
 
 _all_dtypes = (
     int8,
@@ -212,7 +225,7 @@ def __hash__(self):
 }
 
 
-def _result_type(type1, type2):
+def _result_type(type1: DType, type2: DType) -> DType:
     if (type1, type2) in _promotion_table:
         return _promotion_table[type1, type2]
     raise TypeError(f"{type1} and {type2} cannot be type promoted together")
diff --git a/array_api_strict/_elementwise_functions.py b/array_api_strict/_elementwise_functions.py
index c11b17c..6b52a58 100644
--- a/array_api_strict/_elementwise_functions.py
+++ b/array_api_strict/_elementwise_functions.py
@@ -1,51 +1,50 @@
 from __future__ import annotations
 
+import numpy as np
+
+from ._array_object import Array
+from ._creation_functions import asarray
+from ._data_type_functions import broadcast_to, iinfo
 from ._dtypes import (
     _boolean_dtypes,
-    _floating_dtypes,
-    _real_floating_dtypes,
     _complex_floating_dtypes,
+    _dtype_categories,
+    _floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
-    _real_numeric_dtypes,
     _numeric_dtypes,
+    _real_floating_dtypes,
+    _real_numeric_dtypes,
     _result_type,
-    _dtype_categories,
 )
-from ._array_object import Array
 from ._flags import requires_api_version
-from ._creation_functions import asarray
-from ._data_type_functions import broadcast_to, iinfo
 from ._helpers import _maybe_normalize_py_scalars
 
-from typing import Optional, Union
-
-import numpy as np
-
 
 def _binary_ufunc_proto(x1, x2, dtype_category, func_name, np_func):
     """Base implementation of a binary function, `func_name`, defined for
-       dtypes from `dtype_category`
+    dtypes from `dtype_category`
     """
     x1, x2 = _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name)
 
     if x1.device != x2.device:
-        raise ValueError(f"Arrays from two different devices ({x1.device} and {x2.device}) can not be combined.")
+        raise ValueError(
+            f"Arrays from two different devices ({x1.device} and {x2.device}) can not be combined."
+        )
     # Call result type here just to raise on disallowed type combinations
     _result_type(x1.dtype, x2.dtype)
     x1, x2 = Array._normalize_two_args(x1, x2)
     return Array._new(np_func(x1._array, x2._array), device=x1.device)
 
 
-_binary_docstring_template=\
-"""
+_binary_docstring_template = """
 Array API compatible wrapper for :py:func:`np.%s <numpy.%s>`.
 
 See its docstring for more information.
 """
 
 
-def create_binary_func(func_name, dtype_category, np_func):
+def _create_binary_func(func_name, dtype_category, np_func):
     def inner(x1, x2, /) -> Array:
         return _binary_ufunc_proto(x1, x2, dtype_category, func_name, np_func)
     return inner
@@ -58,7 +57,7 @@ def inner(x1, x2, /) -> Array:
     "real numeric": "int | float | Array",
     "numeric": "int | float | complex | Array",
     "integer": "int | Array",
-    "integer or boolean": "int | bool | Array",
+    "integer or boolean": "bool | int | Array",
     "boolean": "bool | Array",
     "real floating-point": "float | Array",
     "complex floating-point": "complex | Array",
@@ -75,7 +74,7 @@ def inner(x1, x2, /) -> Array:
     "bitwise_xor": "integer or boolean",
     "_bitwise_left_shift": "integer",  # leading underscore deliberate
     "_bitwise_right_shift": "integer",
-    # XXX: copysign: real fp or numeric? 
+    # XXX: copysign: real fp or numeric?
     "copysign": "real floating-point",
     "divide": "floating-point",
     "equal": "all",
@@ -105,7 +104,7 @@ def inner(x1, x2, /) -> Array:
     "atan2": "arctan2",
     "_bitwise_left_shift": "left_shift",
     "_bitwise_right_shift": "right_shift",
-    "pow": "power"
+    "pow": "power",
 }
 
 
@@ -117,7 +116,7 @@ def inner(x1, x2, /) -> Array:
     numpy_name = _numpy_renames.get(func_name, func_name)
     np_func = getattr(np, numpy_name)
 
-    func = create_binary_func(func_name, dtype_category, np_func)
+    func = _create_binary_func(func_name, dtype_category, np_func)
     func.__name__ = func_name
 
     func.__doc__ = _binary_docstring_template % (numpy_name, numpy_name)
@@ -153,7 +152,7 @@ def bitwise_right_shift(x1: int | Array, x2: int | Array, /) -> Array:
 
 
 # clean up to not pollute the namespace
-del func, create_binary_func
+del func, _create_binary_func
 
 
 def abs(x: Array, /) -> Array:
@@ -271,8 +270,8 @@ def ceil(x: Array, /) -> Array:
 def clip(
     x: Array,
     /,
-    min: Optional[Union[int, float, Array]] = None,
-    max: Optional[Union[int, float, Array]] = None,
+    min: Array | int | float | None = None,
+    max: Array | int | float | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.clip <numpy.clip>`.
@@ -351,6 +350,7 @@ def clip(
 
     def _isscalar(a):
         return isinstance(a, (int, float, type(None)))
+
     min_shape = () if _isscalar(min) else min.shape
     max_shape = () if _isscalar(max) else max.shape
 
@@ -584,6 +584,7 @@ def reciprocal(x: Array, /) -> Array:
         raise TypeError("Only floating-point dtypes are allowed in reciprocal")
     return Array._new(np.reciprocal(x._array), device=x.device)
 
+
 def round(x: Array, /) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.round <numpy.round>`.
diff --git a/array_api_strict/_fft.py b/array_api_strict/_fft.py
index c888826..2998254 100644
--- a/array_api_strict/_fft.py
+++ b/array_api_strict/_fft.py
@@ -1,31 +1,29 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from collections.abc import Sequence
+from typing import Literal
 
-if TYPE_CHECKING:
-    from typing import Union, Optional, Literal
-    from ._typing import Device, Dtype as DType
-    from collections.abc import Sequence
+import numpy as np
 
+from ._array_object import ALL_DEVICES, Array, Device
+from ._data_type_functions import astype
 from ._dtypes import (
+    DType,
+    _complex_floating_dtypes,
     _floating_dtypes,
     _real_floating_dtypes,
-    _complex_floating_dtypes,
-    float32,
     complex64,
+    float32,
 )
-from ._array_object import Array, ALL_DEVICES
-from ._data_type_functions import astype
 from ._flags import requires_extension
 
-import numpy as np
 
 @requires_extension('fft')
 def fft(
     x: Array,
     /,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
@@ -48,7 +46,7 @@ def ifft(
     x: Array,
     /,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
@@ -71,8 +69,8 @@ def fftn(
     x: Array,
     /,
     *,
-    s: Sequence[int] = None,
-    axes: Sequence[int] = None,
+    s: Sequence[int] | None = None,
+    axes: Sequence[int] | None = None,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
     """
@@ -94,8 +92,8 @@ def ifftn(
     x: Array,
     /,
     *,
-    s: Sequence[int] = None,
-    axes: Sequence[int] = None,
+    s: Sequence[int] | None = None,
+    axes: Sequence[int] | None = None,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
     """
@@ -117,7 +115,7 @@ def rfft(
     x: Array,
     /,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
@@ -140,7 +138,7 @@ def irfft(
     x: Array,
     /,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
@@ -163,8 +161,8 @@ def rfftn(
     x: Array,
     /,
     *,
-    s: Sequence[int] = None,
-    axes: Sequence[int] = None,
+    s: Sequence[int] | None = None,
+    axes: Sequence[int] | None = None,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
     """
@@ -186,8 +184,8 @@ def irfftn(
     x: Array,
     /,
     *,
-    s: Sequence[int] = None,
-    axes: Sequence[int] = None,
+    s: Sequence[int] | None = None,
+    axes: Sequence[int] | None = None,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
     """
@@ -209,7 +207,7 @@ def hfft(
     x: Array,
     /,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
@@ -232,7 +230,7 @@ def ihfft(
     x: Array,
     /,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
@@ -256,8 +254,8 @@ def fftfreq(
     /,
     *,
     d: float = 1.0,
-    dtype: Optional[DType] = None,
-    device: Optional[Device] = None
+    dtype: DType | None = None,
+    device: Device | None = None
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.fftfreq <numpy.fft.fftfreq>`.
@@ -280,8 +278,8 @@ def rfftfreq(
     /,
     *,
     d: float = 1.0,
-    dtype: Optional[DType] = None,
-    device: Optional[Device] = None
+    dtype: DType | None = None,
+    device: Device | None = None
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.rfftfreq <numpy.fft.rfftfreq>`.
@@ -299,7 +297,7 @@ def rfftfreq(
     return Array._new(np_result, device=device)
 
 @requires_extension('fft')
-def fftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
+def fftshift(x: Array, /, *, axes: int | Sequence[int] | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.fftshift <numpy.fft.fftshift>`.
 
@@ -310,7 +308,7 @@ def fftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
     return Array._new(np.fft.fftshift(x._array, axes=axes), device=x.device)
 
 @requires_extension('fft')
-def ifftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
+def ifftshift(x: Array, /, *, axes: int | Sequence[int] | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.ifftshift <numpy.fft.ifftshift>`.
 
diff --git a/array_api_strict/_flags.py b/array_api_strict/_flags.py
index 3fce8a0..6729a4b 100644
--- a/array_api_strict/_flags.py
+++ b/array_api_strict/_flags.py
@@ -12,17 +12,32 @@
 
 """
 
+from __future__ import annotations
+
 import functools
 import os
 import warnings
+from collections.abc import Callable
+from types import TracebackType
+from typing import TYPE_CHECKING, Any, Collection, TypeVar, cast
 
 import array_api_strict
 
+if TYPE_CHECKING:
+    # TODO import from typing (requires Python >= 3.10)
+    from typing_extensions import ParamSpec
+
+    P = ParamSpec("P")
+
+T = TypeVar("T")
+_CallableT = TypeVar("_CallableT", bound=Callable[..., object])
+
+
 supported_versions = (
     "2021.12",
     "2022.12",
     "2023.12",
-    "2024.12"
+    "2024.12",
 )
 
 draft_version = "2025.12"
@@ -43,19 +58,23 @@
     "fft": "2022.12",
 }
 
-ENABLED_EXTENSIONS = default_extensions = (
+default_extensions: tuple[str, ...] = (
     "linalg",
     "fft",
 )
+ENABLED_EXTENSIONS = default_extensions
+
+
 # Public functions
 
+
 def set_array_api_strict_flags(
     *,
-    api_version=None,
-    boolean_indexing=None,
-    data_dependent_shapes=None,
-    enabled_extensions=None,
-):
+    api_version: str | None = None,
+    boolean_indexing: bool | None = None,
+    data_dependent_shapes: bool | None = None,
+    enabled_extensions: Collection[str] | None = None,
+) -> None:
     """
     Set the array-api-strict flags to the specified values.
 
@@ -178,7 +197,8 @@ def set_array_api_strict_flags(
         draft_version=draft_version,
     )
 
-def get_array_api_strict_flags():
+
+def get_array_api_strict_flags() -> dict[str, Any]:
     """
     Get the current array-api-strict flags.
 
@@ -228,7 +248,7 @@ def get_array_api_strict_flags():
     }
 
 
-def reset_array_api_strict_flags():
+def reset_array_api_strict_flags() -> None:
     """
     Reset the array-api-strict flags to their default values.
 
@@ -300,8 +320,19 @@ class ArrayAPIStrictFlags:
     reset_array_api_strict_flags: Reset the flags to their default values.
 
     """
-    def __init__(self, *, api_version=None, boolean_indexing=None,
-                 data_dependent_shapes=None, enabled_extensions=None):
+
+    kwargs: dict[str, Any]
+    old_flags: dict[str, Any]
+    __slots__ = ("kwargs", "old_flags")
+
+    def __init__(
+        self,
+        *,
+        api_version: str | None = None,
+        boolean_indexing: bool | None = None,
+        data_dependent_shapes: bool | None = None,
+        enabled_extensions: Collection[str] | None = None,
+    ):
         self.kwargs = {
             "api_version": api_version,
             "boolean_indexing": boolean_indexing,
@@ -310,12 +341,19 @@ def __init__(self, *, api_version=None, boolean_indexing=None,
         }
         self.old_flags = get_array_api_strict_flags()
 
-    def __enter__(self):
+    def __enter__(self) -> None:
         set_array_api_strict_flags(**self.kwargs)
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+        /,
+    ) -> None:
         set_array_api_strict_flags(**self.old_flags)
 
+
 # Private functions
 
 ENVIRONMENT_VARIABLES = [
@@ -325,8 +363,9 @@ def __exit__(self, exc_type, exc_value, traceback):
     "ARRAY_API_STRICT_ENABLED_EXTENSIONS",
 ]
 
-def set_flags_from_environment():
-    kwargs = {}
+
+def set_flags_from_environment() -> None:
+    kwargs: dict[str, Any] = {}
     if "ARRAY_API_STRICT_API_VERSION" in os.environ:
         kwargs["api_version"] = os.environ["ARRAY_API_STRICT_API_VERSION"]
 
@@ -346,35 +385,41 @@ def set_flags_from_environment():
     # linalg and fft to __all__
     set_array_api_strict_flags(**kwargs)
 
+
 set_flags_from_environment()
 
 # Decorators
 
-def requires_api_version(version):
-    def decorator(func):
+
+def requires_api_version(version: str) -> Callable[[_CallableT], _CallableT]:
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
         @functools.wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             if version > API_VERSION:
                 raise RuntimeError(
                     f"The function {func.__name__} requires API version {version} or later, "
                     f"but the current API version for array-api-strict is {API_VERSION}"
                 )
             return func(*args, **kwargs)
+
         return wrapper
-    return decorator
 
-def requires_data_dependent_shapes(func):
+    return cast(Callable[[_CallableT], _CallableT], decorator)
+
+
+def requires_data_dependent_shapes(func: Callable[P, T]) -> Callable[P, T]:
     @functools.wraps(func)
-    def wrapper(*args, **kwargs):
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
         if not DATA_DEPENDENT_SHAPES:
             raise RuntimeError(f"The function {func.__name__} requires data-dependent shapes, but the data_dependent_shapes flag has been disabled for array-api-strict")
         return func(*args, **kwargs)
     return wrapper
 
-def requires_extension(extension):
-    def decorator(func):
+
+def requires_extension(extension: str) -> Callable[[_CallableT], _CallableT]:
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
         @functools.wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             if extension not in ENABLED_EXTENSIONS:
                 if extension == 'linalg' \
                    and func.__name__ in ['matmul', 'tensordot',
@@ -382,5 +427,7 @@ def wrapper(*args, **kwargs):
                     raise RuntimeError(f"The linalg extension has been disabled for array-api-strict. However, {func.__name__} is also present in the main array_api_strict namespace and may be used from there.")
                 raise RuntimeError(f"The function {func.__name__} requires the {extension} extension, but it has been disabled for array-api-strict")
             return func(*args, **kwargs)
+
         return wrapper
-    return decorator
+
+    return cast(Callable[[_CallableT], _CallableT], decorator)
diff --git a/array_api_strict/_helpers.py b/array_api_strict/_helpers.py
index d3fc9c9..291082e 100644
--- a/array_api_strict/_helpers.py
+++ b/array_api_strict/_helpers.py
@@ -1,18 +1,24 @@
-"""Private helper routines.
-"""
+"""Private helper routines."""
 
-from ._flags import get_array_api_strict_flags
+from __future__ import annotations
+
+from ._array_object import Array
 from ._dtypes import _dtype_categories
+from ._flags import get_array_api_strict_flags
 
 _py_scalars = (bool, int, float, complex)
 
 
-def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
-
+def _maybe_normalize_py_scalars(
+    x1: Array | bool | int | float | complex,
+    x2: Array | bool | int | float | complex,
+    dtype_category: str,
+    func_name: str,
+) -> tuple[Array, Array]:
     flags = get_array_api_strict_flags()
     if flags["api_version"] < "2024.12":
         # scalars will fail at the call site
-        return x1, x2
+        return x1, x2  # type: ignore[return-value]
 
     _allowed_dtypes = _dtype_categories[dtype_category]
 
@@ -34,4 +40,3 @@ def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
             raise TypeError(f"Only {dtype_category} dtypes are allowed in {func_name}(...). "
                             f"Got {x1.dtype} and {x2.dtype}.")
     return x1, x2
-
diff --git a/array_api_strict/_indexing_functions.py b/array_api_strict/_indexing_functions.py
index d7a400e..ab25fab 100644
--- a/array_api_strict/_indexing_functions.py
+++ b/array_api_strict/_indexing_functions.py
@@ -1,17 +1,13 @@
 from __future__ import annotations
 
+import numpy as np
+
 from ._array_object import Array
 from ._dtypes import _integer_dtypes
 from ._flags import requires_api_version
 
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from typing import Optional
 
-import numpy as np
-
-def take(x: Array, indices: Array, /, *, axis: Optional[int] = None) -> Array:
+def take(x: Array, indices: Array, /, *, axis: int | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.take <numpy.take>`.
 
@@ -27,6 +23,7 @@ def take(x: Array, indices: Array, /, *, axis: Optional[int] = None) -> Array:
         raise ValueError(f"Arrays from two different devices ({x.device} and {indices.device}) can not be combined.")
     return Array._new(np.take(x._array, indices._array, axis=axis), device=x.device)
 
+
 @requires_api_version('2024.12')
 def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array:
     """
diff --git a/array_api_strict/_info.py b/array_api_strict/_info.py
index a9dbebf..81f88bc 100644
--- a/array_api_strict/_info.py
+++ b/array_api_strict/_info.py
@@ -1,25 +1,22 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 import numpy as np
 
-if TYPE_CHECKING:
-    from typing import Optional, Union, Tuple, List
-    from ._typing import device, DefaultDataTypes, DataTypes, Capabilities
-
-from ._array_object import ALL_DEVICES, CPU_DEVICE
+from . import _dtypes as dt
+from ._array_object import ALL_DEVICES, CPU_DEVICE, Device
 from ._flags import get_array_api_strict_flags, requires_api_version
-from ._dtypes import bool, int8, int16, int32, int64, uint8, uint16, uint32, uint64, float32, float64, complex64, complex128
+from ._typing import Capabilities, DataTypes, DefaultDataTypes
+
 
 @requires_api_version('2023.12')
 class __array_namespace_info__:
     @requires_api_version('2023.12')
     def capabilities(self) -> Capabilities:
         flags = get_array_api_strict_flags()
-        res = {"boolean indexing": flags['boolean_indexing'],
-                "data-dependent shapes": flags['data_dependent_shapes'],
-                }
+        res: Capabilities = {  # type: ignore[typeddict-item]
+            "boolean indexing": flags['boolean_indexing'],
+            "data-dependent shapes": flags['data_dependent_shapes'],
+        }
         if flags['api_version'] >= '2024.12':
             # maxdims is 32 for NumPy 1.x and 64 for NumPy 2.0. Eventually we will
             # drop support for NumPy 1 but for now, just compute the number
@@ -36,104 +33,104 @@ def capabilities(self) -> Capabilities:
         return res
 
     @requires_api_version('2023.12')
-    def default_device(self) -> device:
+    def default_device(self) -> Device:
         return CPU_DEVICE
 
     @requires_api_version('2023.12')
     def default_dtypes(
         self,
         *,
-        device: Optional[device] = None,
+        device: Device | None = None,
     ) -> DefaultDataTypes:
         return {
-            "real floating": float64,
-            "complex floating": complex128,
-            "integral": int64,
-            "indexing": int64,
+            "real floating": dt.float64,
+            "complex floating": dt.complex128,
+            "integral": dt.int64,
+            "indexing": dt.int64,
         }
 
     @requires_api_version('2023.12')
     def dtypes(
         self,
         *,
-        device: Optional[device] = None,
-        kind: Optional[Union[str, Tuple[str, ...]]] = None,
+        device: Device | None = None,
+        kind: str | tuple[str, ...] | None = None,
     ) -> DataTypes:
         if kind is None:
             return {
-                "bool": bool,
-                "int8": int8,
-                "int16": int16,
-                "int32": int32,
-                "int64": int64,
-                "uint8": uint8,
-                "uint16": uint16,
-                "uint32": uint32,
-                "uint64": uint64,
-                "float32": float32,
-                "float64": float64,
-                "complex64": complex64,
-                "complex128": complex128,
+                "bool": dt.bool,
+                "int8": dt.int8,
+                "int16": dt.int16,
+                "int32": dt.int32,
+                "int64": dt.int64,
+                "uint8": dt.uint8,
+                "uint16": dt.uint16,
+                "uint32": dt.uint32,
+                "uint64": dt.uint64,
+                "float32": dt.float32,
+                "float64": dt.float64,
+                "complex64": dt.complex64,
+                "complex128": dt.complex128,
             }
         if kind == "bool":
-            return {"bool": bool}
+            return {"bool": dt.bool}
         if kind == "signed integer":
             return {
-                "int8": int8,
-                "int16": int16,
-                "int32": int32,
-                "int64": int64,
+                "int8": dt.int8,
+                "int16": dt.int16,
+                "int32": dt.int32,
+                "int64": dt.int64,
             }
         if kind == "unsigned integer":
             return {
-                "uint8": uint8,
-                "uint16": uint16,
-                "uint32": uint32,
-                "uint64": uint64,
+                "uint8": dt.uint8,
+                "uint16": dt.uint16,
+                "uint32": dt.uint32,
+                "uint64": dt.uint64,
             }
         if kind == "integral":
             return {
-                "int8": int8,
-                "int16": int16,
-                "int32": int32,
-                "int64": int64,
-                "uint8": uint8,
-                "uint16": uint16,
-                "uint32": uint32,
-                "uint64": uint64,
+                "int8": dt.int8,
+                "int16": dt.int16,
+                "int32": dt.int32,
+                "int64": dt.int64,
+                "uint8": dt.uint8,
+                "uint16": dt.uint16,
+                "uint32": dt.uint32,
+                "uint64": dt.uint64,
             }
         if kind == "real floating":
             return {
-                "float32": float32,
-                "float64": float64,
+                "float32": dt.float32,
+                "float64": dt.float64,
             }
         if kind == "complex floating":
             return {
-                "complex64": complex64,
-                "complex128": complex128,
+                "complex64": dt.complex64,
+                "complex128": dt.complex128,
             }
         if kind == "numeric":
             return {
-                "int8": int8,
-                "int16": int16,
-                "int32": int32,
-                "int64": int64,
-                "uint8": uint8,
-                "uint16": uint16,
-                "uint32": uint32,
-                "uint64": uint64,
-                "float32": float32,
-                "float64": float64,
-                "complex64": complex64,
-                "complex128": complex128,
+                "int8": dt.int8,
+                "int16": dt.int16,
+                "int32": dt.int32,
+                "int64": dt.int64,
+                "uint8": dt.uint8,
+                "uint16": dt.uint16,
+                "uint32": dt.uint32,
+                "uint64": dt.uint64,
+                "float32": dt.float32,
+                "float64": dt.float64,
+                "complex64": dt.complex64,
+                "complex128": dt.complex128,
             }
         if isinstance(kind, tuple):
-            res = {}
+            res: DataTypes = {}
             for k in kind:
                 res.update(self.dtypes(kind=k))
             return res
         raise ValueError(f"unsupported kind: {kind!r}")
 
     @requires_api_version('2023.12')
-    def devices(self) -> List[device]:
+    def devices(self) -> list[Device]:
         return list(ALL_DEVICES)
diff --git a/array_api_strict/_linalg.py b/array_api_strict/_linalg.py
index 7d379a0..27a2ddf 100644
--- a/array_api_strict/_linalg.py
+++ b/array_api_strict/_linalg.py
@@ -1,33 +1,25 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import partial
+from typing import Literal, NamedTuple
 
-from ._dtypes import (
-    _floating_dtypes,
-    _numeric_dtypes,
-    float32,
-    complex64,
-    complex128,
-)
+import numpy as np
+import numpy.linalg
+
+from ._array_object import Array
 from ._data_type_functions import finfo
-from ._manipulation_functions import reshape
+from ._dtypes import DType, _floating_dtypes, _numeric_dtypes, complex64, complex128
 from ._elementwise_functions import conj
-from ._array_object import Array
-from ._flags import requires_extension, get_array_api_strict_flags
+from ._flags import get_array_api_strict_flags, requires_extension
+from ._manipulation_functions import reshape
+from ._statistical_functions import _np_dtype_sumprod
 
 try:
-    from numpy._core.numeric import normalize_axis_tuple
+    from numpy._core.numeric import normalize_axis_tuple  # type: ignore[attr-defined]
 except ImportError:
-    from numpy.core.numeric import normalize_axis_tuple
+    from numpy.core.numeric import normalize_axis_tuple  # type: ignore[no-redef]
 
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from ._typing import Literal, Optional, Sequence, Tuple, Union, Dtype
-
-from typing import NamedTuple
-
-import numpy.linalg
-import numpy as np
 
 class EighResult(NamedTuple):
     eigenvalues: Array
@@ -175,7 +167,13 @@ def inv(x: Array, /) -> Array:
 # -np.inf, 'fro', 'nuc']]], but Literal does not support floating-point
 # literals.
 @requires_extension('linalg')
-def matrix_norm(x: Array, /, *, keepdims: bool = False, ord: Optional[Union[int, float, Literal['fro', 'nuc']]] = 'fro') -> Array:  # noqa: F821
+def matrix_norm(
+    x: Array,
+    /,
+    *,
+    keepdims: bool = False,
+    ord: float | Literal["fro", "nuc"] | None = "fro",
+) -> Array:  # noqa: F821
     """
     Array API compatible wrapper for :py:func:`np.linalg.norm <numpy.linalg.norm>`.
 
@@ -186,7 +184,10 @@ def matrix_norm(x: Array, /, *, keepdims: bool = False, ord: Optional[Union[int,
     if x.dtype not in _floating_dtypes:
         raise TypeError('Only floating-point dtypes are allowed in matrix_norm')
 
-    return Array._new(np.linalg.norm(x._array, axis=(-2, -1), keepdims=keepdims, ord=ord), device=x.device)
+    return Array._new(
+        np.linalg.norm(x._array, axis=(-2, -1), keepdims=keepdims, ord=ord),
+        device=x.device,
+    )
 
 
 @requires_extension('linalg')
@@ -206,7 +207,7 @@ def matrix_power(x: Array, n: int, /) -> Array:
 
 # Note: the keyword argument name rtol is different from np.linalg.matrix_rank
 @requires_extension('linalg')
-def matrix_rank(x: Array, /, *, rtol: Optional[Union[float, Array]] = None) -> Array:
+def matrix_rank(x: Array, /, *, rtol: float | Array | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.matrix_rank <numpy.matrix_rank>`.
 
@@ -218,13 +219,12 @@ def matrix_rank(x: Array, /, *, rtol: Optional[Union[float, Array]] = None) -> A
         raise np.linalg.LinAlgError("1-dimensional array given. Array must be at least two-dimensional")
     S = np.linalg.svd(x._array, compute_uv=False)
     if rtol is None:
-        tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * finfo(S.dtype).eps
+        tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * np.finfo(S.dtype).eps
     else:
-        if isinstance(rtol, Array):
-            rtol = rtol._array
+        rtol_np = rtol._array if isinstance(rtol, Array) else np.asarray(rtol)
         # Note: this is different from np.linalg.matrix_rank, which does not multiply
         # the tolerance by the largest singular value.
-        tol = S.max(axis=-1, keepdims=True)*np.asarray(rtol)[..., np.newaxis]
+        tol = S.max(axis=-1, keepdims=True) * rtol_np[..., np.newaxis]
     return Array._new(np.count_nonzero(S > tol, axis=-1), device=x.device)
 
 
@@ -252,7 +252,7 @@ def outer(x1: Array, x2: Array, /) -> Array:
 
 # Note: the keyword argument name rtol is different from np.linalg.pinv
 @requires_extension('linalg')
-def pinv(x: Array, /, *, rtol: Optional[Union[float, Array]] = None) -> Array:
+def pinv(x: Array, /, *, rtol: float | Array | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.linalg.pinv <numpy.linalg.pinv>`.
 
@@ -267,9 +267,8 @@ def pinv(x: Array, /, *, rtol: Optional[Union[float, Array]] = None) -> Array:
     # default tolerance by max(M, N).
     if rtol is None:
         rtol = max(x.shape[-2:]) * finfo(x.dtype).eps
-    if isinstance(rtol, Array):
-        rtol = rtol._array
-    return Array._new(np.linalg.pinv(x._array, rcond=rtol), device=x.device)
+    rtol_np = rtol._array if isinstance(rtol, Array) else rtol
+    return Array._new(np.linalg.pinv(x._array, rcond=rtol_np), device=x.device)
 
 @requires_extension('linalg')
 def qr(x: Array, /, *, mode: Literal['reduced', 'complete'] = 'reduced') -> QRResult:  # noqa: F821
@@ -312,14 +311,14 @@ def slogdet(x: Array, /) -> SlogdetResult:
 
 # To workaround this, the below is the code from np.linalg.solve except
 # only calling solve1 in the exactly 1D case.
-def _solve(a, b):
+def _solve(a: np.ndarray, b: np.ndarray) -> np.ndarray:
     try:
-        from numpy.linalg._linalg import (
+        from numpy.linalg._linalg import (  # type: ignore[attr-defined]
         _makearray, _assert_stacked_2d, _assert_stacked_square,
         _commonType, isComplexType, _raise_linalgerror_singular
         )
     except ImportError:
-        from numpy.linalg.linalg import (
+        from numpy.linalg.linalg import (  # type: ignore[attr-defined]
         _makearray, _assert_stacked_2d, _assert_stacked_square,
         _commonType, isComplexType, _raise_linalgerror_singular
         )
@@ -382,14 +381,14 @@ def svd(x: Array, /, *, full_matrices: bool = True) -> SVDResult:
 # Note: svdvals is not in NumPy (but it is in SciPy). It is equivalent to
 # np.linalg.svd(compute_uv=False).
 @requires_extension('linalg')
-def svdvals(x: Array, /) -> Union[Array, Tuple[Array, ...]]:
+def svdvals(x: Array, /) -> Array:
     if x.dtype not in _floating_dtypes:
         raise TypeError('Only floating-point dtypes are allowed in svdvals')
     return Array._new(np.linalg.svd(x._array, compute_uv=False), device=x.device)
 
 # Note: trace is the numpy top-level namespace, not np.linalg
 @requires_extension('linalg')
-def trace(x: Array, /, *, offset: int = 0, dtype: Optional[Dtype] = None) -> Array:
+def trace(x: Array, /, *, offset: int = 0, dtype: DType | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.trace <numpy.trace>`.
 
@@ -398,19 +397,13 @@ def trace(x: Array, /, *, offset: int = 0, dtype: Optional[Dtype] = None) -> Arr
     if x.dtype not in _numeric_dtypes:
         raise TypeError('Only numeric dtypes are allowed in trace')
 
-    # Note: trace() works the same as sum() and prod() (see
-    # _statistical_functions.py)
-    if dtype is None:
-        if get_array_api_strict_flags()['api_version'] < '2023.12':
-            if x.dtype == float32:
-                dtype = np.float64
-            elif x.dtype == complex64:
-                dtype = np.complex128
-    else:
-        dtype = dtype._np_dtype
+    # Note: trace() works the same as sum() and prod() (see _statistical_functions.py)
+    np_dtype = _np_dtype_sumprod(x, dtype)
+
     # Note: trace always operates on the last two axes, whereas np.trace
     # operates on the first two axes by default
-    return Array._new(np.asarray(np.trace(x._array, offset=offset, axis1=-2, axis2=-1, dtype=dtype)), device=x.device)
+    res = np.trace(x._array, offset=offset, axis1=-2, axis2=-1, dtype=np_dtype)
+    return Array._new(np.asarray(res), device=x.device)
 
 # Note: the name here is different from norm(). The array API norm is split
 # into matrix_norm and vector_norm().
@@ -418,7 +411,14 @@ def trace(x: Array, /, *, offset: int = 0, dtype: Optional[Dtype] = None) -> Arr
 # The type for ord should be Optional[Union[int, float, Literal[np.inf,
 # -np.inf]]] but Literal does not support floating-point literals.
 @requires_extension('linalg')
-def vector_norm(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False, ord: Optional[Union[int, float]] = 2) -> Array:
+def vector_norm(
+    x: Array,
+    /,
+    *,
+    axis: int | tuple[int, ...] | None = None,
+    keepdims: bool = False,
+    ord: int | float = 2,
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.linalg.norm <numpy.linalg.norm>`.
 
@@ -456,8 +456,8 @@ def vector_norm(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = No
         # We can't reuse np.linalg.norm(keepdims) because of the reshape hacks
         # above to avoid matrix norm logic.
         shape = list(x.shape)
-        _axis = normalize_axis_tuple(range(x.ndim) if axis is None else axis, x.ndim)
-        for i in _axis:
+        axis_tup = normalize_axis_tuple(range(x.ndim) if axis is None else axis, x.ndim)
+        for i in axis_tup:
             shape[i] = 1
         res = reshape(res, tuple(shape))
 
@@ -480,7 +480,13 @@ def matmul(x1: Array, x2: Array, /) -> Array:
 
 # Note: tensordot is the numpy top-level namespace but not in np.linalg
 @requires_extension('linalg')
-def tensordot(x1: Array, x2: Array, /, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2) -> Array:
+def tensordot(
+    x1: Array,
+    x2: Array,
+    /,
+    *,
+    axes: int | tuple[Sequence[int], Sequence[int]] = 2,
+) -> Array:
     from ._linear_algebra_functions import tensordot
     return tensordot(x1, x2, axes=axes)
 
diff --git a/array_api_strict/_linear_algebra_functions.py b/array_api_strict/_linear_algebra_functions.py
index 6af2a15..d18214c 100644
--- a/array_api_strict/_linear_algebra_functions.py
+++ b/array_api_strict/_linear_algebra_functions.py
@@ -7,16 +7,15 @@
 
 from __future__ import annotations
 
-from ._dtypes import _numeric_dtypes
+from collections.abc import Sequence
+
+import numpy as np
+import numpy.linalg
+
 from ._array_object import Array
+from ._dtypes import _numeric_dtypes
 from ._flags import get_array_api_strict_flags
 
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from ._typing import Sequence, Tuple, Union
-
-import numpy.linalg
-import numpy as np
 
 # Note: matmul is the numpy top-level namespace but not in np.linalg
 def matmul(x1: Array, x2: Array, /) -> Array:
@@ -38,7 +37,13 @@ def matmul(x1: Array, x2: Array, /) -> Array:
 # Note: tensordot is the numpy top-level namespace but not in np.linalg
 
 # Note: axes must be a tuple, unlike np.tensordot where it can be an array or array-like.
-def tensordot(x1: Array, x2: Array, /, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2) -> Array:
+def tensordot(
+    x1: Array,
+    x2: Array,
+    /,
+    *,
+    axes: int | tuple[Sequence[int], Sequence[int]] = 2,
+) -> Array:
     # Note: the restriction to numeric dtypes only is different from
     # np.tensordot.
     if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
diff --git a/array_api_strict/_manipulation_functions.py b/array_api_strict/_manipulation_functions.py
index 63c3516..e2fd24c 100644
--- a/array_api_strict/_manipulation_functions.py
+++ b/array_api_strict/_manipulation_functions.py
@@ -1,21 +1,17 @@
 from __future__ import annotations
 
+import numpy as np
+
 from ._array_object import Array
 from ._creation_functions import asarray
 from ._data_type_functions import astype, result_type
 from ._dtypes import _integer_dtypes, int64, uint64
-from ._flags import requires_api_version, get_array_api_strict_flags
-
-from typing import TYPE_CHECKING
+from ._flags import get_array_api_strict_flags, requires_api_version
 
-if TYPE_CHECKING:
-    from typing import List, Optional, Tuple, Union
-
-import numpy as np
 
 # Note: the function name is different here
 def concat(
-    arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: Optional[int] = 0
+    arrays: tuple[Array, ...] | list[Array], /, *, axis: int | None = 0
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.concatenate <numpy.concatenate>`.
@@ -29,8 +25,11 @@ def concat(
         raise ValueError("concat inputs must all be on the same device")
     result_device = arrays[0].device
 
-    arrays = tuple(a._array for a in arrays)
-    return Array._new(np.concatenate(arrays, axis=axis, dtype=dtype._np_dtype), device=result_device)
+    np_arrays = tuple(a._array for a in arrays)
+    return Array._new(
+        np.concatenate(np_arrays, axis=axis, dtype=dtype._np_dtype),
+        device=result_device,
+    )
 
 
 def expand_dims(x: Array, /, *, axis: int) -> Array:
@@ -42,7 +41,7 @@ def expand_dims(x: Array, /, *, axis: int) -> Array:
     return Array._new(np.expand_dims(x._array, axis), device=x.device)
 
 
-def flip(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> Array:
+def flip(x: Array, /, *, axis: int | tuple[int, ...] | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.flip <numpy.flip>`.
 
@@ -53,8 +52,8 @@ def flip(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) ->
 @requires_api_version('2023.12')
 def moveaxis(
     x: Array,
-    source: Union[int, Tuple[int, ...]],
-    destination: Union[int, Tuple[int, ...]],
+    source: int | tuple[int, ...],
+    destination: int | tuple[int, ...],
     /,
 ) -> Array:
     """
@@ -66,7 +65,7 @@ def moveaxis(
 
 # Note: The function name is different here (see also matrix_transpose).
 # Unlike transpose(), the axes argument is required.
-def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array:
+def permute_dims(x: Array, /, axes: tuple[int, ...]) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.transpose <numpy.transpose>`.
 
@@ -77,10 +76,10 @@ def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array:
 @requires_api_version('2023.12')
 def repeat(
     x: Array,
-    repeats: Union[int, Array],
+    repeats: int | Array,
     /,
     *,
-    axis: Optional[int] = None,
+    axis: int | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.repeat <numpy.repeat>`.
@@ -108,12 +107,9 @@ def repeat(
         repeats = astype(repeats, int64)
     return Array._new(np.repeat(x._array, repeats._array, axis=axis), device=x.device)
 
+
 # Note: the optional argument is called 'shape', not 'newshape'
-def reshape(x: Array,
-            /,
-            shape: Tuple[int, ...],
-            *,
-            copy: Optional[bool] = None) -> Array:
+def reshape(x: Array, /, shape: tuple[int, ...], *, copy: bool | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.reshape <numpy.reshape>`.
 
@@ -135,9 +131,9 @@ def reshape(x: Array,
 def roll(
     x: Array,
     /,
-    shift: Union[int, Tuple[int, ...]],
+    shift: int | tuple[int, ...],
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.roll <numpy.roll>`.
@@ -147,7 +143,7 @@ def roll(
     return Array._new(np.roll(x._array, shift, axis=axis), device=x.device)
 
 
-def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array:
+def squeeze(x: Array, /, axis: int | tuple[int, ...]) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.squeeze <numpy.squeeze>`.
 
@@ -161,7 +157,7 @@ def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array:
     return Array._new(np.squeeze(x._array, axis=axis), device=x.device)
 
 
-def stack(arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: int = 0) -> Array:
+def stack(arrays: tuple[Array, ...] | list[Array], /, *, axis: int = 0) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.stack <numpy.stack>`.
 
@@ -172,12 +168,12 @@ def stack(arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: int = 0) ->
     if len({a.device for a in arrays}) > 1:
         raise ValueError("concat inputs must all be on the same device")
     result_device = arrays[0].device
-    arrays = tuple(a._array for a in arrays)
-    return Array._new(np.stack(arrays, axis=axis), device=result_device)
+    np_arrays = tuple(a._array for a in arrays)
+    return Array._new(np.stack(np_arrays, axis=axis), device=result_device)
 
 
 @requires_api_version('2023.12')
-def tile(x: Array, repetitions: Tuple[int, ...], /) -> Array:
+def tile(x: Array, repetitions: tuple[int, ...], /) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.tile <numpy.tile>`.
 
@@ -190,7 +186,7 @@ def tile(x: Array, repetitions: Tuple[int, ...], /) -> Array:
 
 # Note: this function is new
 @requires_api_version('2023.12')
-def unstack(x: Array, /, *, axis: int = 0) -> Tuple[Array, ...]:
+def unstack(x: Array, /, *, axis: int = 0) -> tuple[Array, ...]:
     if not (-x.ndim <= axis < x.ndim):
         raise ValueError("axis out of range")
 
diff --git a/array_api_strict/_searching_functions.py b/array_api_strict/_searching_functions.py
index 9864132..b366ed9 100644
--- a/array_api_strict/_searching_functions.py
+++ b/array_api_strict/_searching_functions.py
@@ -1,18 +1,17 @@
 from __future__ import annotations
 
-from ._array_object import Array
-from ._dtypes import _result_type, _real_numeric_dtypes, bool as _bool
-from ._flags import requires_data_dependent_shapes, requires_api_version, get_array_api_strict_flags
-from ._helpers import _maybe_normalize_py_scalars
-
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from typing import Literal, Optional, Tuple, Union
+from typing import Literal
 
 import numpy as np
 
+from ._array_object import Array
+from ._dtypes import _real_numeric_dtypes, _result_type
+from ._dtypes import bool as _bool
+from ._flags import requires_api_version, requires_data_dependent_shapes
+from ._helpers import _maybe_normalize_py_scalars
+
 
-def argmax(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -> Array:
+def argmax(x: Array, /, *, axis: int | None = None, keepdims: bool = False) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.argmax <numpy.argmax>`.
 
@@ -23,7 +22,7 @@ def argmax(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -
     return Array._new(np.asarray(np.argmax(x._array, axis=axis, keepdims=keepdims)), device=x.device)
 
 
-def argmin(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -> Array:
+def argmin(x: Array, /, *, axis: int | None = None, keepdims: bool = False) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.argmin <numpy.argmin>`.
 
@@ -35,7 +34,7 @@ def argmin(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -
 
 
 @requires_data_dependent_shapes
-def nonzero(x: Array, /) -> Tuple[Array, ...]:
+def nonzero(x: Array, /) -> tuple[Array, ...]:
     """
     Array API compatible wrapper for :py:func:`np.nonzero <numpy.nonzero>`.
 
@@ -52,7 +51,7 @@ def count_nonzero(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     """
@@ -71,7 +70,7 @@ def searchsorted(
     /,
     *,
     side: Literal["left", "right"] = "left",
-    sorter: Optional[Array] = None,
+    sorter: Array | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.searchsorted <numpy.searchsorted>`.
@@ -84,25 +83,29 @@ def searchsorted(
     if x1.device != x2.device:
         raise ValueError(f"Arrays from two different devices ({x1.device} and {x2.device}) can not be combined.")
 
-    sorter = sorter._array if sorter is not None else None
+    np_sorter = sorter._array if sorter is not None else None
     # TODO: The sort order of nans and signed zeros is implementation
     # dependent. Should we error/warn if they are present?
 
     # x1 must be 1-D, but NumPy already requires this.
-    return Array._new(np.searchsorted(x1._array, x2._array, side=side, sorter=sorter), device=x1.device)
+    return Array._new(
+        np.searchsorted(x1._array, x2._array, side=side, sorter=np_sorter), 
+        device=x1.device,
+    )
+
 
 def where(
     condition: Array,
-    x1: bool | int | float | complex | Array,
-    x2: bool | int | float | complex | Array, /
+    x1: Array | bool | int | float | complex,
+    x2: Array | bool | int | float | complex,
+    /,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.where <numpy.where>`.
 
     See its docstring for more information.
     """
-    if get_array_api_strict_flags()['api_version'] > '2023.12':
-        x1, x2 = _maybe_normalize_py_scalars(x1, x2, "all", "where")
+    x1, x2 = _maybe_normalize_py_scalars(x1, x2, "all", "where")
 
     # Call result type here just to raise on disallowed type combinations
     _result_type(x1.dtype, x2.dtype)
diff --git a/array_api_strict/_set_functions.py b/array_api_strict/_set_functions.py
index 7bd5bad..e677a52 100644
--- a/array_api_strict/_set_functions.py
+++ b/array_api_strict/_set_functions.py
@@ -1,13 +1,12 @@
 from __future__ import annotations
 
-from ._array_object import Array
-
-from ._flags import requires_data_dependent_shapes
-
 from typing import NamedTuple
 
 import numpy as np
 
+from ._array_object import Array
+from ._flags import requires_data_dependent_shapes
+
 # Note: np.unique() is split into four functions in the array API:
 # unique_all, unique_counts, unique_inverse, and unique_values (this is done
 # to remove polymorphic return types).
@@ -20,6 +19,7 @@
 # Note: The functions here return a namedtuple (np.unique() returns a normal
 # tuple).
 
+
 class UniqueAllResult(NamedTuple):
     values: Array
     indices: Array
diff --git a/array_api_strict/_sorting_functions.py b/array_api_strict/_sorting_functions.py
index 765bd9e..e9193f1 100644
--- a/array_api_strict/_sorting_functions.py
+++ b/array_api_strict/_sorting_functions.py
@@ -1,10 +1,12 @@
 from __future__ import annotations
 
-from ._array_object import Array
-from ._dtypes import _real_numeric_dtypes
+from typing import Literal
 
 import numpy as np
 
+from ._array_object import Array
+from ._dtypes import _real_numeric_dtypes
+
 
 # Note: the descending keyword argument is new in this function
 def argsort(
@@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
     if not descending:
         res = np.argsort(x._array, axis=axis, kind=kind)
     else:
@@ -35,6 +37,7 @@ def argsort(
         res = max_i - res
     return Array._new(res, device=x.device)
 
+
 # Note: the descending keyword argument is new in this function
 def sort(
     x: Array, /, *, axis: int = -1, descending: bool = False, stable: bool = True
@@ -47,8 +50,7 @@ def sort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in sort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
-    res = np.sort(x._array, axis=axis, kind=kind)
+    res = np.sort(x._array, axis=axis, kind="stable" if stable else "quicksort")
     if descending:
         res = np.flip(res, axis=axis)
     return Array._new(res, device=x.device)
diff --git a/array_api_strict/_statistical_functions.py b/array_api_strict/_statistical_functions.py
index e41e7ef..668cd02 100644
--- a/array_api_strict/_statistical_functions.py
+++ b/array_api_strict/_statistical_functions.py
@@ -1,38 +1,36 @@
 from __future__ import annotations
 
+from typing import Any
+
+import numpy as np
+
+from ._array_object import Array
+from ._creation_functions import ones, zeros
 from ._dtypes import (
-    _real_floating_dtypes,
-    _real_numeric_dtypes,
+    DType,
     _floating_dtypes,
+    _np_dtype,
     _numeric_dtypes,
+    _real_floating_dtypes,
+    _real_numeric_dtypes,
+    complex64,
+    float32,
 )
-from ._array_object import Array
-from ._dtypes import float32, complex64
-from ._flags import requires_api_version, get_array_api_strict_flags
-from ._creation_functions import zeros, ones
+from ._flags import get_array_api_strict_flags, requires_api_version
 from ._manipulation_functions import concat
 
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union
-    from ._typing import Dtype
-
-import numpy as np
 
 @requires_api_version('2023.12')
 def cumulative_sum(
     x: Array,
     /,
     *,
-    axis: Optional[int] = None,
-    dtype: Optional[Dtype] = None,
+    axis: int | None = None,
+    dtype: DType | None = None,
     include_initial: bool = False,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in cumulative_sum")
-    if dtype is not None:
-        dtype = dtype._np_dtype
 
     # TODO: The standard is not clear about what should happen when x.ndim == 0.
     if axis is None:
@@ -44,7 +42,7 @@ def cumulative_sum(
         if axis < 0:
             axis += x.ndim
         x = concat([zeros(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype), x], axis=axis)
-    return Array._new(np.cumsum(x._array, axis=axis, dtype=dtype), device=x.device)
+    return Array._new(np.cumsum(x._array, axis=axis, dtype=_np_dtype(dtype)), device=x.device)
 
 
 @requires_api_version('2024.12')
@@ -52,8 +50,8 @@ def cumulative_prod(
     x: Array,
     /,
     *,
-    axis: Optional[int] = None,
-    dtype: Optional[Dtype] = None,
+    axis: int | None = None,
+    dtype: DType | None = None,
     include_initial: bool = False,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
@@ -61,9 +59,6 @@ def cumulative_prod(
     if x.ndim == 0:
         raise ValueError("Only ndim >= 1 arrays are allowed in cumulative_prod")
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-
     if axis is None:
         if x.ndim > 1:
             raise ValueError("axis must be specified in cumulative_prod for more than one dimension")
@@ -74,14 +69,14 @@ def cumulative_prod(
         if axis < 0:
             axis += x.ndim
         x = concat([ones(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype), x], axis=axis)
-    return Array._new(np.cumprod(x._array, axis=axis, dtype=dtype), device=x.device)
+    return Array._new(np.cumprod(x._array, axis=axis, dtype=_np_dtype(dtype)), device=x.device)
 
 
 def max(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     if x.dtype not in _real_numeric_dtypes:
@@ -93,14 +88,15 @@ def mean(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
 
-    if get_array_api_strict_flags()['api_version'] > '2023.12':
-        allowed_dtypes = _floating_dtypes
-    else:
-        allowed_dtypes = _real_floating_dtypes
+    allowed_dtypes = (
+        _floating_dtypes
+        if get_array_api_strict_flags()['api_version'] > '2023.12'
+        else _real_floating_dtypes
+    )
 
     if x.dtype not in allowed_dtypes:
         raise TypeError("Only floating-point dtypes are allowed in mean")
@@ -111,7 +107,7 @@ def min(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     if x.dtype not in _real_numeric_dtypes:
@@ -119,37 +115,43 @@ def min(
     return Array._new(np.min(x._array, axis=axis, keepdims=keepdims), device=x.device)
 
 
+def _np_dtype_sumprod(x: Array, dtype: DType | None) -> np.dtype[Any] | None:
+    """In versions prior to 2023.12, sum() and prod() upcast for all
+    dtypes when dtype=None. For 2023.12, the behavior is the same as in
+    NumPy (only upcast for integral dtypes).
+    """
+    if dtype is None and get_array_api_strict_flags()['api_version'] < '2023.12':
+        if x.dtype == float32:
+            return np.float64  # type: ignore[return-value]
+        elif x.dtype == complex64:
+            return np.complex128  # type: ignore[return-value]
+    return _np_dtype(dtype)
+
+
 def prod(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    dtype: Optional[Dtype] = None,
+    axis: int | tuple[int, ...] | None = None,
+    dtype: DType | None = None,
     keepdims: bool = False,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in prod")
 
-    if dtype is None:
-        # Note: In versions prior to 2023.12, sum() and prod() upcast for all
-        # dtypes when dtype=None. For 2023.12, the behavior is the same as in
-        # NumPy (only upcast for integral dtypes).
-        if get_array_api_strict_flags()['api_version'] < '2023.12':
-            if x.dtype == float32:
-                dtype = np.float64
-            elif x.dtype == complex64:
-                dtype = np.complex128
-    else:
-        dtype = dtype._np_dtype
-    return Array._new(np.prod(x._array, dtype=dtype, axis=axis, keepdims=keepdims), device=x.device)
+    np_dtype = _np_dtype_sumprod(x, dtype)
+    return Array._new(
+        np.prod(x._array, dtype=np_dtype, axis=axis, keepdims=keepdims),
+        device=x.device,
+    )
 
 
 def std(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    correction: Union[int, float] = 0.0,
+    axis: int | tuple[int, ...] | None = None,
+    correction: int | float = 0.0,
     keepdims: bool = False,
 ) -> Array:
     # Note: the keyword argument correction is different here
@@ -162,33 +164,26 @@ def sum(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    dtype: Optional[Dtype] = None,
+    axis: int | tuple[int, ...] | None = None,
+    dtype: DType | None = None,
     keepdims: bool = False,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in sum")
 
-    if dtype is None:
-        # Note: In versions prior to 2023.12, sum() and prod() upcast for all
-        # dtypes when dtype=None. For 2023.12, the behavior is the same as in
-        # NumPy (only upcast for integral dtypes).
-        if get_array_api_strict_flags()['api_version'] < '2023.12':
-            if x.dtype == float32:
-                dtype = np.float64
-            elif x.dtype == complex64:
-                dtype = np.complex128
-    else:
-        dtype = dtype._np_dtype
-    return Array._new(np.sum(x._array, axis=axis, dtype=dtype, keepdims=keepdims), device=x.device)
+    np_dtype = _np_dtype_sumprod(x, dtype)
+    return Array._new(
+        np.sum(x._array, axis=axis, dtype=np_dtype, keepdims=keepdims),
+        device=x.device,
+    )
 
 
 def var(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    correction: Union[int, float] = 0.0,
+    axis: int | tuple[int, ...] | None = None,
+    correction: int | float = 0.0,
     keepdims: bool = False,
 ) -> Array:
     # Note: the keyword argument correction is different here
diff --git a/array_api_strict/_typing.py b/array_api_strict/_typing.py
index 94c4975..91095a8 100644
--- a/array_api_strict/_typing.py
+++ b/array_api_strict/_typing.py
@@ -8,41 +8,19 @@
 
 from __future__ import annotations
 
-__all__ = [
-    "Array",
-    "Device",
-    "Dtype",
-    "SupportsDLPack",
-    "SupportsBufferProtocol",
-    "PyCapsule",
-]
-
 import sys
+from typing import Any, Protocol, TypedDict, TypeVar
 
-from typing import (
-    Any,
-    TypedDict,
-    TypeVar,
-    Protocol,
-)
-
-from ._array_object import Array, _device
-from ._dtypes import _DType
-from ._info import __array_namespace_info__
+from ._dtypes import DType
 
 _T_co = TypeVar("_T_co", covariant=True)
 
+
 class NestedSequence(Protocol[_T_co]):
     def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
     def __len__(self, /) -> int: ...
 
 
-Device = _device
-
-Dtype = _DType
-
-Info = __array_namespace_info__
-
 if sys.version_info >= (3, 12):
     from collections.abc import Buffer as SupportsBufferProtocol
 else:
@@ -50,40 +28,42 @@ def __len__(self, /) -> int: ...
 
 PyCapsule = Any
 
+
 class SupportsDLPack(Protocol):
     def __dlpack__(self, /, *, stream: None = ...) -> PyCapsule: ...
 
+
 Capabilities = TypedDict(
-    "Capabilities", {"boolean indexing": bool, "data-dependent shapes": bool,
-                     "max dimensions": int}
+    "Capabilities",
+    {
+        "boolean indexing": bool,
+        "data-dependent shapes": bool,
+        "max dimensions": int,
+    },
 )
 
 DefaultDataTypes = TypedDict(
     "DefaultDataTypes",
     {
-        "real floating": Dtype,
-        "complex floating": Dtype,
-        "integral": Dtype,
-        "indexing": Dtype,
+        "real floating": DType,
+        "complex floating": DType,
+        "integral": DType,
+        "indexing": DType,
     },
 )
 
-DataTypes = TypedDict(
-    "DataTypes",
-    {
-        "bool": Dtype,
-        "float32": Dtype,
-        "float64": Dtype,
-        "complex64": Dtype,
-        "complex128": Dtype,
-        "int8": Dtype,
-        "int16": Dtype,
-        "int32": Dtype,
-        "int64": Dtype,
-        "uint8": Dtype,
-        "uint16": Dtype,
-        "uint32": Dtype,
-        "uint64": Dtype,
-    },
-    total=False,
-)
+
+class DataTypes(TypedDict, total=False):
+    bool: DType
+    float32: DType
+    float64: DType
+    complex64: DType
+    complex128: DType
+    int8: DType
+    int16: DType
+    int32: DType
+    int64: DType
+    uint8: DType
+    uint16: DType
+    uint32: DType
+    uint64: DType
diff --git a/array_api_strict/_utility_functions.py b/array_api_strict/_utility_functions.py
index f75f36f..fab1025 100644
--- a/array_api_strict/_utility_functions.py
+++ b/array_api_strict/_utility_functions.py
@@ -1,21 +1,20 @@
 from __future__ import annotations
 
-from ._array_object import Array
-from ._flags import requires_api_version
-from ._dtypes import _numeric_dtypes
-
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union
+from typing import Any
 
 import numpy as np
+import numpy.typing as npt
+
+from ._array_object import Array
+from ._dtypes import _numeric_dtypes
+from ._flags import requires_api_version
 
 
 def all(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     """
@@ -30,7 +29,7 @@ def any(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     """
@@ -40,6 +39,7 @@ def any(
     """
     return Array._new(np.asarray(np.any(x._array, axis=axis, keepdims=keepdims)), device=x.device)
 
+
 @requires_api_version('2024.12')
 def diff(
     x: Array,
@@ -47,8 +47,8 @@ def diff(
     *,
     axis: int = -1,
     n: int = 1,
-    prepend: Optional[Array] = None,
-    append: Optional[Array] = None,
+    prepend: Array | None = None,
+    append: Array | None = None,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in diff")
@@ -57,7 +57,7 @@ def diff(
     # currently specified.
 
     # NumPy does not support prepend=None or append=None
-    kwargs = dict(axis=axis, n=n)
+    kwargs: dict[str, int | npt.NDArray[Any]] = {"axis": axis, "n": n}
     if prepend is not None:
         if prepend.device != x.device:
             raise ValueError(f"Arrays from two different devices ({prepend.device} and {x.device}) can not be combined.")
diff --git a/array_api_strict/py.typed b/array_api_strict/py.typed
new file mode 100644
index 0000000..e69de29
diff --git a/array_api_strict/tests/test_validation.py b/array_api_strict/tests/test_validation.py
index bd76ec6..5552e3a 100644
--- a/array_api_strict/tests/test_validation.py
+++ b/array_api_strict/tests/test_validation.py
@@ -1,11 +1,9 @@
-from typing import Callable
-
 import pytest
 
 import array_api_strict as xp
 
 
-def p(func: Callable, *args, **kwargs):
+def p(func, *args, **kwargs):
     f_sig = ", ".join(
         [str(a) for a in args] + [f"{k}={v}" for k, v in kwargs.items()]
     )
diff --git a/pyproject.toml b/pyproject.toml
index b3d2594..cf9e6dd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,3 +31,18 @@ Repository = "https://github.com/data-apis/array-api-strict"
 [tool.setuptools_scm]
 version_file = "array_api_strict/_version.py"
 
+[tool.mypy]
+disallow_incomplete_defs = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+show_error_codes = true
+warn_redundant_casts = true
+warn_unused_ignores = true
+warn_unreachable = true
+strict_bytes = true
+local_partial_types = true
+
+[[tool.mypy.overrides]]
+module = ["*.tests.*"]
+disallow_untyped_defs = false


Reviews:
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T15:59:38Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T16:03:05Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T16:05:33Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T16:05:42Z
user: login: ev-br
body: Thanks for the PR Guido!

I went through up to `_linalg.py` before running out of steam for the week. In general it looks great, my main gripe is about annotations for python scalars:

- several look like global search-and-replace : floordiv, `__gt__` just do not accept `complex`
- I'm -1 for changing `bool | int | float` into just `float`.  
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-21T17:09:25Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T17:16:53Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T17:17:43Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T17:22:00Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T17:23:13Z
user: login: NeilGirdhar
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-21T17:23:16Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T17:25:36Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T17:27:16Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T17:29:23Z
user: login: ev-br
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-21T17:45:52Z
user: login: ev-br
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-21T17:50:57Z
user: login: ev-br
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-21T17:51:53Z
user: login: ev-br
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-21T17:56:36Z
user: login: NeilGirdhar
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-21T18:26:27Z
user: login: jorenham
body: Good stuff! I left a couple of suggestions; do with it what you want 👌🏻
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-21T20:38:13Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T21:20:30Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T21:24:07Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T21:24:55Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T21:26:59Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T21:28:31Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T21:30:03Z
user: login: jorenham
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-22T16:32:51Z
user: login: jorenham
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-22T16:37:26Z
user: login: jorenham
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-22T16:40:46Z
user: login: jorenham
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-22T16:44:39Z
user: login: jorenham
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-22T16:49:13Z
user: login: jorenham
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-22T16:50:42Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T14:06:15Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T14:19:09Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T14:22:15Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T14:23:43Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T14:27:36Z
user: login: ev-br
body: Let me record my -1 on `bool | int |float` until https://github.com/data-apis/array-api-compat/pull/257#issuecomment-2746365323 is done.
state: DISMISSED
author_association: MEMBER
submitted_at: 2025-03-26T16:04:44Z
user: login: jorenham
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-26T18:48:56Z
user: login: jorenham
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-26T18:50:15Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T18:58:55Z
user: login: jorenham
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-26T18:59:46Z
user: login: jorenham
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-26T19:19:32Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T13:28:41Z
user: login: jorenham
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-27T13:54:36Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T17:35:28Z
user: login: jorenham
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-28T14:44:00Z
user: login: lucascolley
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-29T22:16:54Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-31T13:22:05Z

Review comments:
diff_hunk: @@ -16,62 +16,66 @@
 from __future__ import annotations
 
 import operator
+from collections.abc import Iterator
 from enum import IntEnum
+from types import ModuleType
+from typing import Any, Literal, SupportsIndex
 
-from ._creation_functions import asarray
+import numpy as np
+import numpy.typing as npt
+
+from ._creation_functions import _default, _Default, asarray
 from ._dtypes import (
-    _DType,
+    DType,
     _all_dtypes,
     _boolean_dtypes,
+    _complex_floating_dtypes,
+    _dtype_categories,
+    _floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
-    _floating_dtypes,
-    _real_floating_dtypes,
-    _complex_floating_dtypes,
     _numeric_dtypes,
-    _result_type,
-    _dtype_categories,
+    _real_floating_dtypes,
     _real_to_complex_map,
+    _result_type,
 )
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
+from ._typing import PyCapsule
 
-from typing import TYPE_CHECKING, SupportsIndex
-import types
-
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union, Any
-    from ._typing import PyCapsule, Dtype
-    import numpy.typing as npt
-
-import numpy as np
+try:
+    from types import EllipsisType  # Python >=3.10
+except ImportError:
+    EllipsisType = type(Ellipsis)
 
 class Device:
-    def __init__(self, device="CPU_DEVICE"):
+    _device: str
+    __slots__ = ("_device", "__weakref__")
path: array_api_strict/_array_object.py
user: login: ev-br
body: What is this for here? AFAIU this forces the lack of `__dict__`, do we want to save memory or prohibit adding attrbitutes or something else? 
created_at: 2025-03-21T15:53:48Z
author_association: MEMBER
diff_hunk: @@ -313,12 +305,14 @@ def linspace(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.linspace(start, stop, num, dtype=dtype, endpoint=endpoint), device=device)
+    return Array._new(
+        np.linspace(start, stop, num, dtype=_np_dtype(dtype), endpoint=endpoint),
+        device=device,
+    )
 
 
-def meshgrid(*arrays: Array, indexing: str = "xy") -> List[Array]:
+# Note: indexing was 'str' in <=2024.12
+def meshgrid(*arrays: Array, indexing: str = "xy") -> list[Array]:
path: array_api_strict/_creation_functions.py
user: login: crusaderky
body: https://github.com/data-apis/array-api/pull/915
created_at: 2025-03-21T15:59:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -16,62 +16,66 @@
 from __future__ import annotations
 
 import operator
+from collections.abc import Iterator
 from enum import IntEnum
+from types import ModuleType
+from typing import Any, Literal, SupportsIndex
 
-from ._creation_functions import asarray
+import numpy as np
+import numpy.typing as npt
+
+from ._creation_functions import _default, _Default, asarray
 from ._dtypes import (
-    _DType,
+    DType,
     _all_dtypes,
     _boolean_dtypes,
+    _complex_floating_dtypes,
+    _dtype_categories,
+    _floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
-    _floating_dtypes,
-    _real_floating_dtypes,
-    _complex_floating_dtypes,
     _numeric_dtypes,
-    _result_type,
-    _dtype_categories,
+    _real_floating_dtypes,
     _real_to_complex_map,
+    _result_type,
 )
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
+from ._typing import PyCapsule
 
-from typing import TYPE_CHECKING, SupportsIndex
-import types
-
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union, Any
-    from ._typing import PyCapsule, Dtype
-    import numpy.typing as npt
-
-import numpy as np
+try:
+    from types import EllipsisType  # Python >=3.10
+except ImportError:
+    EllipsisType = type(Ellipsis)
 
 class Device:
-    def __init__(self, device="CPU_DEVICE"):
+    _device: str
+    __slots__ = ("_device", "__weakref__")
+
+    def __init__(self, device: str = "CPU_DEVICE"):
         if device not in ("CPU_DEVICE", "device1", "device2"):
             raise ValueError(f"The device '{device}' is not a valid choice.")
         self._device = device
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"array_api_strict.Device('{self._device}')"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
path: array_api_strict/_array_object.py
user: login: ev-br
body: Returns an array of the bool dtype
created_at: 2025-03-21T16:02:06Z
author_association: MEMBER
diff_hunk: @@ -313,12 +305,14 @@ def linspace(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.linspace(start, stop, num, dtype=dtype, endpoint=endpoint), device=device)
+    return Array._new(
+        np.linspace(start, stop, num, dtype=_np_dtype(dtype), endpoint=endpoint),
+        device=device,
+    )
 
 
-def meshgrid(*arrays: Array, indexing: str = "xy") -> List[Array]:
+# Note: indexing was 'str' in <=2024.12
+def meshgrid(*arrays: Array, indexing: str = "xy") -> list[Array]:
path: array_api_strict/_creation_functions.py
user: login: crusaderky
body: not sure if this is inconsequential enough that I should just go on and define it as Literal here?
created_at: 2025-03-21T16:03:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -140,34 +138,32 @@ def finfo(type: Union[Dtype, Array], /) -> finfo_object:
         float(fi.max),
         float(fi.min),
         float(fi.smallest_normal),
-        fi.dtype,
+        DType(fi.dtype),
path: array_api_strict/_data_type_functions.py
user: login: crusaderky
body: Fix runtime bug
created_at: 2025-03-21T16:05:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -140,34 +138,32 @@ def finfo(type: Union[Dtype, Array], /) -> finfo_object:
         float(fi.max),
         float(fi.min),
         float(fi.smallest_normal),
-        fi.dtype,
+        DType(fi.dtype),
     )
 
 
-def iinfo(type: Union[Dtype, Array], /) -> iinfo_object:
+def iinfo(type: DType | Array, /) -> iinfo_object:
     """
     Array API compatible wrapper for :py:func:`np.iinfo <numpy.iinfo>`.
 
     See its docstring for more information.
     """
-    if isinstance(type, _DType):
-        type = type._np_dtype
-    ii = np.iinfo(type)
-    return iinfo_object(ii.bits, ii.max, ii.min, ii.dtype)
+    np_type = type._array if isinstance(type, Array) else type._np_dtype
+    ii = np.iinfo(np_type)
+    return iinfo_object(ii.bits, ii.max, ii.min, DType(ii.dtype))
path: array_api_strict/_data_type_functions.py
user: login: crusaderky
body: Fix runtime bug
created_at: 2025-03-21T16:05:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -88,11 +92,14 @@ class Array:
 
     """
     _array: npt.NDArray[Any]
+    _dtype: DType
+    _device: Device
+    __slots__ = ("_array", "_dtype", "_device", "__weakref__")
 
     # Use a custom constructor instead of __init__, as manually initializing
     # this class is not supported API.
     @classmethod
-    def _new(cls, x, /, device):
+    def _new(cls, x: np.ndarray | np.generic, /, device: Device | None) -> Array:
path: array_api_strict/_array_object.py
user: login: ev-br
body: 1. `np.ndarray` here but `_array` above is annotated as `npt.NDarray[Any]`. These two should be consistent, I'd think?

2. `Device | None` hints at the default value being None, but it is required. To double-check my understanding, this annotation means that the matching runtime check is
`if device is not None or not isinstance(device, Device): raise TypeError()`, correct? 
created_at: 2025-03-21T16:23:22Z
author_association: MEMBER
diff_hunk: @@ -120,15 +127,15 @@ def _new(cls, x, /, device):
         return obj
 
     # Prevent Array() from working
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *args: object, **kwargs: object) -> Array:
path: array_api_strict/_array_object.py
user: login: ev-br
body: do we really need the `*args: object` annotation, what does it buy us
created_at: 2025-03-21T16:24:10Z
author_association: MEMBER
diff_hunk: @@ -184,7 +193,9 @@ def __array__(self, dtype: None | np.dtype[Any] = None, copy: None | bool = None
     # spec in places where it either deviates from or is more strict than
     # NumPy behavior
 
-    def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_category: str, op: str) -> Array:
+    def _check_allowed_dtypes(
+        self, other: Array | complex, dtype_category: str, op: str
path: array_api_strict/_array_object.py
user: login: ev-br
body: ints, bools and floats are allowed. I get that a float is-a complex but nonetheless. As long as this is for a human reader more than for a linter, let's be more descriptive. Here and elsewhere.
 
created_at: 2025-03-21T16:27:16Z
author_association: MEMBER
diff_hunk: @@ -225,16 +236,18 @@ def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_categor
 
         return other
 
-    def _check_device(self, other):
+    def _check_device(self, other: Array | complex) -> None:
path: array_api_strict/_array_object.py
user: login: ev-br
body: as above
created_at: 2025-03-21T16:28:19Z
author_association: MEMBER
diff_hunk: @@ -225,16 +236,18 @@ def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_categor
 
         return other
 
-    def _check_device(self, other):
+    def _check_device(self, other: Array | complex) -> None:
         """Check that other is on a device compatible with the current array"""
         if isinstance(other, (int, complex, float, bool)):
             return
         elif isinstance(other, Array):
             if self.device != other.device:
                 raise ValueError(f"Arrays from two different devices ({self.device} and {other.device}) can not be combined.")
+        else:
+            raise TypeError(f"Expected Array | python scalar; got {type(other)}")
 
     # Helper function to match the type promotion rules in the spec
-    def _promote_scalar(self, scalar):
+    def _promote_scalar(self, scalar: complex) -> Array:
path: array_api_strict/_array_object.py
user: login: ev-br
body: as above
created_at: 2025-03-21T16:28:28Z
author_association: MEMBER
diff_hunk: @@ -518,7 +541,7 @@ def __abs__(self: Array, /) -> Array:
         res = self._array.__abs__()
         return self.__class__._new(res, device=self.device)
 
-    def __add__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __add__(self, other: Array | complex, /) -> Array:
path: array_api_strict/_array_object.py
user: login: ev-br
body: as above. I'm going to stop flagging specific instances, please roll them back everywhere.
(changing `Union[x, y]` to `x | y` is great, however)
created_at: 2025-03-21T16:30:04Z
author_association: MEMBER
diff_hunk: @@ -655,7 +676,7 @@ def __float__(self: Array, /) -> float:
         res = self._array.__float__()
         return res
 
-    def __floordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __floordiv__(self, other: Array | complex, /) -> Array:
path: array_api_strict/_array_object.py
user: login: ev-br
body: Here it's definitely wrong: `__floordiv__` does not accept complex arguments:

```
In [6]: xp.arange(3) // 1j
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[6], line 1
...
TypeError: Python complex scalars can only be promoted with floating-point arrays.

```
created_at: 2025-03-21T16:36:05Z
author_association: MEMBER
diff_hunk: @@ -667,7 +688,7 @@ def __floordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__floordiv__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ge__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __ge__(self, other: Array | complex, /) -> Array:
path: array_api_strict/_array_object.py
user: login: ev-br
body: ditto : comparison is undefined in the complex plane.

(Am going to stop flagging individual instances.)
created_at: 2025-03-21T16:36:54Z
author_association: MEMBER
diff_hunk: @@ -680,14 +701,14 @@ def __ge__(self: Array, other: Union[int, float, Array], /) -> Array:
         return self.__class__._new(res, device=self.device)
 
     def __getitem__(
-        self: Array,
-        key: Union[
-            int,
-            slice,
-            ellipsis,  # noqa: F821
-            Tuple[Union[int, slice, ellipsis, None], ...],  # noqa: F821
-            Array,
-        ],
+        self,
+        key: (
+            int
+            | slice
+            | EllipsisType
+            | tuple[int | slice | EllipsisType | None, ...]
+            | Array
+        ),
path: array_api_strict/_array_object.py
user: login: ev-br
body: Strictly speaking, tuples of arrays are fine. But to express it in full we'll need way way more (overloads?), so OK to punt on this.
created_at: 2025-03-21T16:38:07Z
author_association: MEMBER
diff_hunk: @@ -715,7 +735,7 @@ def __gt__(self: Array, other: Union[int, float, Array], /) -> Array:
         res = self._array.__gt__(other._array)
         return self.__class__._new(res, device=other.device)
 
-    def __int__(self: Array, /) -> int:
+    def __int__(self) -> int:
path: array_api_strict/_array_object.py
user: login: ev-br
body: So the rule is that you never explicitly type annotate `self`? That makes sense indeed.
created_at: 2025-03-21T16:41:19Z
author_association: MEMBER
diff_hunk: @@ -136,16 +139,17 @@ def arange(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.arange(start, stop=stop, step=step, dtype=dtype), device=device)
+    return Array._new(
+        np.arange(start, stop, step, dtype=_np_dtype(dtype)), 
+        device=device,
+    )
path: array_api_strict/_creation_functions.py
user: login: ev-br
body: `<rant>` it's a bit annoying that for whatever reason the layout changes here (linters?) but not in empty(...) below. Was easier to read when it was consistent. `</rant>`
created_at: 2025-03-21T16:46:27Z
author_association: MEMBER
diff_hunk: @@ -201,19 +201,15 @@ def eye(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.eye(n_rows, M=n_cols, k=k, dtype=dtype), device=device)
+    return Array._new(np.eye(n_rows, M=n_cols, k=k, dtype=_np_dtype(dtype)), device=device)
 
 
-_default = object()
-
 def from_dlpack(
-    x: object,
+    x: SupportsDLPack,
     /,
     *,
-    device: Optional[Device] = _default,
-    copy: Optional[bool] = _default,
+    device: Device | _Default | None = _default,
path: array_api_strict/_creation_functions.py
user: login: ev-br
body: So it's now OK to show users something that starts with a leading underscore? `_Device` here
created_at: 2025-03-21T16:49:34Z
author_association: MEMBER
diff_hunk: @@ -1,19 +1,27 @@
+from __future__ import annotations
path: array_api_strict/_dtypes.py
user: login: ev-br
body: I thought this is no longer needed?
created_at: 2025-03-21T16:51:46Z
author_association: MEMBER
diff_hunk: @@ -1,19 +1,27 @@
+from __future__ import annotations
+
 import warnings
+from typing import Any
 
 import numpy as np
+import numpy.typing as npt
 
 # Note: we wrap the NumPy dtype objects in a bare class, so that none of the
 # additional methods and behaviors of NumPy dtype objects are exposed.
 
-class _DType:
-    def __init__(self, np_dtype):
-        np_dtype = np.dtype(np_dtype)
-        self._np_dtype = np_dtype
+py_bool = bool
+
+class DType:
+    _np_dtype: np.dtype[Any]
+    __slots__ = ("_np_dtype", "__weakref__")
+
+    def __init__(self, np_dtype: npt.DTypeLike):
+        self._np_dtype = np.dtype(np_dtype)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"array_api_strict.{self._np_dtype.name}"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> py_bool:
path: array_api_strict/_dtypes.py
user: login: ev-br
body: builtins.bool?
created_at: 2025-03-21T16:52:47Z
author_association: MEMBER
diff_hunk: @@ -218,13 +225,12 @@ def matrix_rank(x: Array, /, *, rtol: Optional[Union[float, Array]] = None) -> A
         raise np.linalg.LinAlgError("1-dimensional array given. Array must be at least two-dimensional")
     S = np.linalg.svd(x._array, compute_uv=False)
     if rtol is None:
-        tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * finfo(S.dtype).eps
+        tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * np.finfo(S.dtype).eps
path: array_api_strict/_linalg.py
user: login: ev-br
body: Note to self: seems OK, if only to save an import.
created_at: 2025-03-21T16:59:51Z
author_association: MEMBER
diff_hunk: @@ -16,62 +16,67 @@
 from __future__ import annotations
 
 import operator
+from collections.abc import Iterator
 from enum import IntEnum
+from types import ModuleType
+from typing import Any, Literal, SupportsIndex
 
-from ._creation_functions import asarray
+import numpy as np
+import numpy.typing as npt
+
+from ._creation_functions import _default, _Default, asarray
 from ._dtypes import (
-    _DType,
+    DType,
     _all_dtypes,
     _boolean_dtypes,
+    _complex_floating_dtypes,
+    _dtype_categories,
+    _floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
-    _floating_dtypes,
-    _real_floating_dtypes,
-    _complex_floating_dtypes,
     _numeric_dtypes,
-    _result_type,
-    _dtype_categories,
+    _real_floating_dtypes,
     _real_to_complex_map,
+    _result_type,
 )
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
+from ._typing import PyCapsule
 
-from typing import TYPE_CHECKING, SupportsIndex
-import types
+try:
+    from types import EllipsisType  # Python >=3.10
+except ImportError:
+    EllipsisType = type(Ellipsis)
 
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union, Any
-    from ._typing import PyCapsule, Dtype
-    import numpy.typing as npt
-
-import numpy as np
 
 class Device:
-    def __init__(self, device="CPU_DEVICE"):
+    _device: str
path: array_api_strict/_array_object.py
user: login: jorenham
body: since it's hashable:

```suggestion
    _device: Final[str]
```
created_at: 2025-03-21T17:06:50Z
author_association: MEMBER
diff_hunk: @@ -16,62 +16,67 @@
 from __future__ import annotations
 
 import operator
+from collections.abc import Iterator
 from enum import IntEnum
+from types import ModuleType
+from typing import Any, Literal, SupportsIndex
 
-from ._creation_functions import asarray
+import numpy as np
+import numpy.typing as npt
+
+from ._creation_functions import _default, _Default, asarray
 from ._dtypes import (
-    _DType,
+    DType,
     _all_dtypes,
     _boolean_dtypes,
+    _complex_floating_dtypes,
+    _dtype_categories,
+    _floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
-    _floating_dtypes,
-    _real_floating_dtypes,
-    _complex_floating_dtypes,
     _numeric_dtypes,
-    _result_type,
-    _dtype_categories,
+    _real_floating_dtypes,
     _real_to_complex_map,
+    _result_type,
 )
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
+from ._typing import PyCapsule
 
-from typing import TYPE_CHECKING, SupportsIndex
-import types
+try:
+    from types import EllipsisType  # Python >=3.10
+except ImportError:
+    EllipsisType = type(Ellipsis)
path: array_api_strict/_array_object.py
user: login: jorenham
body: Type-checkers don't understand `try:` blocks like these, but they *do* understand `if sys.version_info >= (3, 10):`. 

The `type(Ellipsis)` also won't be understood by them. You could work around that on python 3.9 by doing something like

```py
if sys.version_info >= (3, 10):
    from typing import EllipsisType
else:
    if TYPE_CHECKING:
        from typing_extensions import EllipsisType
    else:
        EllipsisType = type(Ellipsis)
```
created_at: 2025-03-21T17:11:24Z
author_association: MEMBER
diff_hunk: @@ -16,62 +16,67 @@
 from __future__ import annotations
 
 import operator
+from collections.abc import Iterator
 from enum import IntEnum
+from types import ModuleType
+from typing import Any, Literal, SupportsIndex
 
-from ._creation_functions import asarray
+import numpy as np
+import numpy.typing as npt
+
+from ._creation_functions import _default, _Default, asarray
 from ._dtypes import (
-    _DType,
+    DType,
     _all_dtypes,
     _boolean_dtypes,
+    _complex_floating_dtypes,
+    _dtype_categories,
+    _floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
-    _floating_dtypes,
-    _real_floating_dtypes,
-    _complex_floating_dtypes,
     _numeric_dtypes,
-    _result_type,
-    _dtype_categories,
+    _real_floating_dtypes,
     _real_to_complex_map,
+    _result_type,
 )
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
+from ._typing import PyCapsule
 
-from typing import TYPE_CHECKING, SupportsIndex
-import types
+try:
+    from types import EllipsisType  # Python >=3.10
+except ImportError:
+    EllipsisType = type(Ellipsis)
 
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union, Any
-    from ._typing import PyCapsule, Dtype
-    import numpy.typing as npt
-
-import numpy as np
 
 class Device:
-    def __init__(self, device="CPU_DEVICE"):
+    _device: str
path: array_api_strict/_array_object.py
user: login: jorenham
body: Perhaps the entire `Device` could be made `@final` too 🤔 
created_at: 2025-03-21T17:12:50Z
author_association: MEMBER
diff_hunk: @@ -16,62 +16,66 @@
 from __future__ import annotations
 
 import operator
+from collections.abc import Iterator
 from enum import IntEnum
+from types import ModuleType
+from typing import Any, Literal, SupportsIndex
 
-from ._creation_functions import asarray
+import numpy as np
+import numpy.typing as npt
+
+from ._creation_functions import _default, _Default, asarray
 from ._dtypes import (
-    _DType,
+    DType,
     _all_dtypes,
     _boolean_dtypes,
+    _complex_floating_dtypes,
+    _dtype_categories,
+    _floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
-    _floating_dtypes,
-    _real_floating_dtypes,
-    _complex_floating_dtypes,
     _numeric_dtypes,
-    _result_type,
-    _dtype_categories,
+    _real_floating_dtypes,
     _real_to_complex_map,
+    _result_type,
 )
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
+from ._typing import PyCapsule
 
-from typing import TYPE_CHECKING, SupportsIndex
-import types
-
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union, Any
-    from ._typing import PyCapsule, Dtype
-    import numpy.typing as npt
-
-import numpy as np
+try:
+    from types import EllipsisType  # Python >=3.10
+except ImportError:
+    EllipsisType = type(Ellipsis)
 
 class Device:
-    def __init__(self, device="CPU_DEVICE"):
+    _device: str
+    __slots__ = ("_device", "__weakref__")
path: array_api_strict/_array_object.py
user: login: crusaderky
body: It's a performance tweak. I can revert if you feel strongly against it.
created_at: 2025-03-21T17:16:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -16,62 +16,66 @@
 from __future__ import annotations
 
 import operator
+from collections.abc import Iterator
 from enum import IntEnum
+from types import ModuleType
+from typing import Any, Literal, SupportsIndex
 
-from ._creation_functions import asarray
+import numpy as np
+import numpy.typing as npt
+
+from ._creation_functions import _default, _Default, asarray
 from ._dtypes import (
-    _DType,
+    DType,
     _all_dtypes,
     _boolean_dtypes,
+    _complex_floating_dtypes,
+    _dtype_categories,
+    _floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
-    _floating_dtypes,
-    _real_floating_dtypes,
-    _complex_floating_dtypes,
     _numeric_dtypes,
-    _result_type,
-    _dtype_categories,
+    _real_floating_dtypes,
     _real_to_complex_map,
+    _result_type,
 )
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
+from ._typing import PyCapsule
 
-from typing import TYPE_CHECKING, SupportsIndex
-import types
-
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union, Any
-    from ._typing import PyCapsule, Dtype
-    import numpy.typing as npt
-
-import numpy as np
+try:
+    from types import EllipsisType  # Python >=3.10
+except ImportError:
+    EllipsisType = type(Ellipsis)
 
 class Device:
-    def __init__(self, device="CPU_DEVICE"):
+    _device: str
+    __slots__ = ("_device", "__weakref__")
+
+    def __init__(self, device: str = "CPU_DEVICE"):
         if device not in ("CPU_DEVICE", "device1", "device2"):
             raise ValueError(f"The device '{device}' is not a valid choice.")
         self._device = device
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"array_api_strict.Device('{self._device}')"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
path: array_api_strict/_array_object.py
user: login: crusaderky
body: this is `Device.__eq__`, not `Array.__eq__`
created_at: 2025-03-21T17:17:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -88,11 +92,14 @@ class Array:
 
     """
     _array: npt.NDArray[Any]
+    _dtype: DType
+    _device: Device
+    __slots__ = ("_array", "_dtype", "_device", "__weakref__")
 
     # Use a custom constructor instead of __init__, as manually initializing
     # this class is not supported API.
     @classmethod
-    def _new(cls, x, /, device):
+    def _new(cls, x: np.ndarray | np.generic, /, device: Device | None) -> Array:
path: array_api_strict/_array_object.py
user: login: crusaderky
body: The implementation explicitly caters for None or omitted. `asarray` leverages this use case e.g. `xp.asarray(1)`
created_at: 2025-03-21T17:22:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -120,15 +127,15 @@ def _new(cls, x, /, device):
         return obj
 
     # Prevent Array() from working
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *args: object, **kwargs: object) -> Array:
path: array_api_strict/_array_object.py
user: login: crusaderky
body: It buys us being able to run mypy with
```
disallow_incomplete_defs = true
disallow_untyped_defs = true
```
created_at: 2025-03-21T17:23:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -184,7 +193,9 @@ def __array__(self, dtype: None | np.dtype[Any] = None, copy: None | bool = None
     # spec in places where it either deviates from or is more strict than
     # NumPy behavior
 
-    def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_category: str, op: str) -> Array:
+    def _check_allowed_dtypes(
+        self, other: Array | complex, dtype_category: str, op: str
path: array_api_strict/_array_object.py
user: login: NeilGirdhar
body: > As long as this is for a human reader more than for a linter

I think @jorenham is right that this has a lot of drawbacks.  Once you get used to annotations, you know that complex subsumes all of those options.  If you want to make things clear for novices, why not use a docstring?
created_at: 2025-03-21T17:23:16Z
author_association: NONE
diff_hunk: @@ -184,7 +193,9 @@ def __array__(self, dtype: None | np.dtype[Any] = None, copy: None | bool = None
     # spec in places where it either deviates from or is more strict than
     # NumPy behavior
 
-    def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_category: str, op: str) -> Array:
+    def _check_allowed_dtypes(
+        self, other: Array | complex, dtype_category: str, op: str
path: array_api_strict/_array_object.py
user: login: crusaderky
body: Here's the point though, none of this is for a human reader that isn't pushing a PR to this repo, save for the public functions in `_flags.py`!

https://data-apis.org/array-api-strict/api.html
created_at: 2025-03-21T17:25:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -201,19 +201,15 @@ def eye(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.eye(n_rows, M=n_cols, k=k, dtype=dtype), device=device)
+    return Array._new(np.eye(n_rows, M=n_cols, k=k, dtype=_np_dtype(dtype)), device=device)
 
 
-_default = object()
-
 def from_dlpack(
-    x: object,
+    x: SupportsDLPack,
     /,
     *,
-    device: Optional[Device] = _default,
-    copy: Optional[bool] = _default,
+    device: Device | _Default | None = _default,
path: array_api_strict/_creation_functions.py
user: login: crusaderky
body: this isn't rendered. However, I'll rename it to `Undef` for clarity.
created_at: 2025-03-21T17:27:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,19 +1,27 @@
+from __future__ import annotations
path: array_api_strict/_dtypes.py
user: login: crusaderky
body: This PR is still for python 3.9.
In Python >=3.10, you'll still need it whenever there is a type annotation before the definition, or a conditional import inside `if TYPE_CHECKING:`.
created_at: 2025-03-21T17:29:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -120,15 +127,15 @@ def _new(cls, x, /, device):
         return obj
 
     # Prevent Array() from working
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *args: object, **kwargs: object) -> Array:
path: array_api_strict/_array_object.py
user: login: ev-br
body: Fine by me, as long as the additional strictness is optional. 
created_at: 2025-03-21T17:45:52Z
author_association: MEMBER
diff_hunk: @@ -16,62 +16,66 @@
 from __future__ import annotations
 
 import operator
+from collections.abc import Iterator
 from enum import IntEnum
+from types import ModuleType
+from typing import Any, Literal, SupportsIndex
 
-from ._creation_functions import asarray
+import numpy as np
+import numpy.typing as npt
+
+from ._creation_functions import _default, _Default, asarray
 from ._dtypes import (
-    _DType,
+    DType,
     _all_dtypes,
     _boolean_dtypes,
+    _complex_floating_dtypes,
+    _dtype_categories,
+    _floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
-    _floating_dtypes,
-    _real_floating_dtypes,
-    _complex_floating_dtypes,
     _numeric_dtypes,
-    _result_type,
-    _dtype_categories,
+    _real_floating_dtypes,
     _real_to_complex_map,
+    _result_type,
 )
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
+from ._typing import PyCapsule
 
-from typing import TYPE_CHECKING, SupportsIndex
-import types
-
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union, Any
-    from ._typing import PyCapsule, Dtype
-    import numpy.typing as npt
-
-import numpy as np
+try:
+    from types import EllipsisType  # Python >=3.10
+except ImportError:
+    EllipsisType = type(Ellipsis)
 
 class Device:
-    def __init__(self, device="CPU_DEVICE"):
+    _device: str
+    __slots__ = ("_device", "__weakref__")
path: array_api_strict/_array_object.py
user: login: ev-br
body: No strong feelings here. IME slots as a perf or memory optimization are not worth it. But fine to add them here. If a need for a dunder dict appears (numpy considers adding it to funds, FWIW), we'll reconsider. 
created_at: 2025-03-21T17:50:57Z
author_association: MEMBER
diff_hunk: @@ -88,11 +92,14 @@ class Array:
 
     """
     _array: npt.NDArray[Any]
+    _dtype: DType
+    _device: Device
+    __slots__ = ("_array", "_dtype", "_device", "__weakref__")
 
     # Use a custom constructor instead of __init__, as manually initializing
     # this class is not supported API.
     @classmethod
-    def _new(cls, x, /, device):
+    def _new(cls, x: np.ndarray | np.generic, /, device: Device | None) -> Array:
path: array_api_strict/_array_object.py
user: login: ev-br
body: 
Thanks. The annotation question still stands though. 
 EDIT: resolved in the last commit.
created_at: 2025-03-21T17:51:53Z
author_association: MEMBER
diff_hunk: @@ -184,7 +193,9 @@ def __array__(self, dtype: None | np.dtype[Any] = None, copy: None | bool = None
     # spec in places where it either deviates from or is more strict than
     # NumPy behavior
 
-    def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_category: str, op: str) -> Array:
+    def _check_allowed_dtypes(
+        self, other: Array | complex, dtype_category: str, op: str
path: array_api_strict/_array_object.py
user: login: ev-br
body: Sorry, no. If the tooling is not ready, we don't use it until it is. Please roll it back, here and in the -strict PR.
created_at: 2025-03-21T17:56:36Z
author_association: MEMBER
diff_hunk: @@ -184,7 +193,9 @@ def __array__(self, dtype: None | np.dtype[Any] = None, copy: None | bool = None
     # spec in places where it either deviates from or is more strict than
     # NumPy behavior
 
-    def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_category: str, op: str) -> Array:
+    def _check_allowed_dtypes(
+        self, other: Array | complex, dtype_category: str, op: str
path: array_api_strict/_array_object.py
user: login: NeilGirdhar
body: What do you mean by the "the tooling is not ready"?  Python typing will probably never alter the definition of these basic types, right?
created_at: 2025-03-21T18:26:27Z
author_association: NONE
diff_hunk: @@ -184,7 +193,9 @@ def __array__(self, dtype: None | np.dtype[Any] = None, copy: None | bool = None
     # spec in places where it either deviates from or is more strict than
     # NumPy behavior
 
-    def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_category: str, op: str) -> Array:
+    def _check_allowed_dtypes(
+        self, other: Array | complex, dtype_category: str, op: str
path: array_api_strict/_array_object.py
user: login: jorenham
body: > What do you mean by the "the tooling is not ready"? Python typing will probably never alter the definition of these basic types, right?

That's right. 
created_at: 2025-03-21T19:26:47Z
author_association: MEMBER
diff_hunk: @@ -807,7 +829,7 @@ def __matmul__(self: Array, other: Array, /) -> Array:
         res = self._array.__matmul__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __mod__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __mod__(self, other: Array | complex, /) -> Array:
path: array_api_strict/_array_object.py
user: login: jorenham
body: I think this should be

```suggestion
    def __mod__(self, other: Array | float, /) -> Array:
```
created_at: 2025-03-21T19:30:42Z
author_association: MEMBER
diff_hunk: @@ -1013,7 +1034,7 @@ def __ifloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
         self._array.__ifloordiv__(other._array)
         return self
 
-    def __rfloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __rfloordiv__(self, other: Array | complex, /) -> Array:
path: array_api_strict/_array_object.py
user: login: jorenham
body: ```suggestion
    def __rfloordiv__(self, other: Array | float, /) -> Array:
```
created_at: 2025-03-21T19:32:52Z
author_association: MEMBER
diff_hunk: @@ -1002,7 +1023,7 @@ def __rand__(self: Array, other: Union[int, bool, Array], /) -> Array:
         res = self._array.__rand__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ifloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __ifloordiv__(self, other: Array | complex, /) -> Array:
path: array_api_strict/_array_object.py
user: login: jorenham
body: ```suggestion
    def __ifloordiv__(self, other: Array | float, /) -> Array:
```
created_at: 2025-03-21T19:33:00Z
author_association: MEMBER
diff_hunk: @@ -1084,7 +1105,7 @@ def __imod__(self: Array, other: Union[int, float, Array], /) -> Array:
         self._array.__imod__(other._array)
         return self
 
-    def __rmod__(self: Array, other: Union[int, float, Array], /) -> Array:
+    def __rmod__(self, other: Array | complex, /) -> Array:
path: array_api_strict/_array_object.py
user: login: jorenham
body: ```suggestion
    def __rmod__(self, other: Array | float, /) -> Array:
```
created_at: 2025-03-21T19:34:01Z
author_association: MEMBER
diff_hunk: @@ -1,23 +1,30 @@
 from __future__ import annotations
 
+from collections.abc import Generator
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+from enum import Enum
+from typing import TYPE_CHECKING, Literal, cast
 
-if TYPE_CHECKING:
-    from ._typing import (
-        Array,
-        Device,
-        Dtype,
-        NestedSequence,
-        SupportsBufferProtocol,
-    )
-from ._dtypes import _DType, _all_dtypes
+import numpy as np
+
+from ._dtypes import DType, _all_dtypes, _np_dtype
 from ._flags import get_array_api_strict_flags
+from ._typing import NestedSequence, SupportsBufferProtocol, SupportsDLPack
+
+if TYPE_CHECKING:
+    # Circular import
+    from ._array_object import Array, Device
+
+
+class Undef(Enum):
+    UNDEF = 0
+
+
+_undef = Undef.UNDEF
 
-import numpy as np
 
 @contextmanager
-def allow_array():
+def allow_array() -> Generator[None, None, None]:
path: array_api_strict/_creation_functions.py
user: login: jorenham
body: Maybe it's only a stub thing, but I believe that this is also allowed nowadays

```suggestion
def allow_array() -> Generator[None]:
```
created_at: 2025-03-21T19:36:03Z
author_association: MEMBER
diff_hunk: @@ -30,43 +37,40 @@ def allow_array():
     finally:
         _array_object._allow_array = original_value
 
-def _check_valid_dtype(dtype):
+
+def _check_valid_dtype(dtype: DType | None) -> None:
     # Note: Only spelling dtypes as the dtype objects is supported.
     if dtype not in (None,) + _all_dtypes:
         raise ValueError(f"dtype must be one of the supported dtypes, got {dtype!r}")
 
-def _supports_buffer_protocol(obj):
+
+def _supports_buffer_protocol(obj: object) -> bool:
path: array_api_strict/_creation_functions.py
user: login: jorenham
body: How about something like:

```suggestion
def _supports_buffer_protocol(obj: object) -> TypeGuard[SupportsBufferProtocol]:
```

or better yet, use `TypeIs` from `typing_extensions`
created_at: 2025-03-21T19:38:12Z
author_association: MEMBER
diff_hunk: @@ -310,12 +340,18 @@ def __init__(self, *, api_version=None, boolean_indexing=None,
         }
         self.old_flags = get_array_api_strict_flags()
 
-    def __enter__(self):
+    def __enter__(self) -> None:
         set_array_api_strict_flags(**self.kwargs)
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
path: array_api_strict/_flags.py
user: login: jorenham
body: ```suggestion
        traceback: TracebackType | None,
        /,
```
created_at: 2025-03-21T19:45:35Z
author_association: MEMBER
diff_hunk: @@ -346,35 +383,41 @@ def set_flags_from_environment():
     # linalg and fft to __all__
     set_array_api_strict_flags(**kwargs)
 
+
 set_flags_from_environment()
 
 # Decorators
 
-def requires_api_version(version):
-    def decorator(func):
+
+def requires_api_version(version: str) -> Callable[[Callable], Callable]:
path: array_api_strict/_flags.py
user: login: jorenham
body: How about something like this?

```suggestion
_CallableT = TypeVar("_CallableT", bound=Callable[..., object])

def requires_api_version(version: str) -> Callable[[_CallableT], _CallableT]:
```

This way, the decorator doesn't delete the signature of the decoratee
created_at: 2025-03-21T19:50:22Z
author_association: MEMBER
diff_hunk: @@ -346,35 +383,41 @@ def set_flags_from_environment():
     # linalg and fft to __all__
     set_array_api_strict_flags(**kwargs)
 
+
 set_flags_from_environment()
 
 # Decorators
 
-def requires_api_version(version):
-    def decorator(func):
+
+def requires_api_version(version: str) -> Callable[[Callable], Callable]:
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
         @functools.wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             if version > API_VERSION:
                 raise RuntimeError(
                     f"The function {func.__name__} requires API version {version} or later, "
                     f"but the current API version for array-api-strict is {API_VERSION}"
                 )
             return func(*args, **kwargs)
+
         return wrapper
+
     return decorator
 
-def requires_data_dependent_shapes(func):
+
+def requires_data_dependent_shapes(func: Callable[P, T]) -> Callable[P, T]:
     @functools.wraps(func)
-    def wrapper(*args, **kwargs):
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
         if not DATA_DEPENDENT_SHAPES:
             raise RuntimeError(f"The function {func.__name__} requires data-dependent shapes, but the data_dependent_shapes flag has been disabled for array-api-strict")
         return func(*args, **kwargs)
     return wrapper
 
-def requires_extension(extension):
-    def decorator(func):
+
+def requires_extension(extension: str) -> Callable[[Callable], Callable]:
path: array_api_strict/_flags.py
user: login: jorenham
body: same opportunity for improvement as above

```suggestion
def requires_extension(extension: str) -> Callable[[_CallableT], _CallableT]:
```
created_at: 2025-03-21T19:51:01Z
author_association: MEMBER
diff_hunk: @@ -1,18 +1,24 @@
-"""Private helper routines.
-"""
+"""Private helper routines."""
 
-from ._flags import get_array_api_strict_flags
+from __future__ import annotations
+
+from ._array_object import Array
 from ._dtypes import _dtype_categories
+from ._flags import get_array_api_strict_flags
 
 _py_scalars = (bool, int, float, complex)
 
 
-def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
-
+def _maybe_normalize_py_scalars(
+    x1: Array | complex,
+    x2: Array | complex,
+    dtype_category: str,
+    func_name: str,
+) -> tuple[Array, Array]:
     flags = get_array_api_strict_flags()
     if flags["api_version"] < "2024.12":
         # scalars will fail at the call site
-        return x1, x2
+        return x1, x2  # type: ignore[return-value]
path: array_api_strict/_helpers.py
user: login: jorenham
body: a `cast` would be a bit cleaner
created_at: 2025-03-21T19:51:50Z
author_association: MEMBER
diff_hunk: @@ -36,104 +33,104 @@ def capabilities(self) -> Capabilities:
         return res
 
     @requires_api_version('2023.12')
-    def default_device(self) -> device:
+    def default_device(self) -> Device:
         return CPU_DEVICE
 
     @requires_api_version('2023.12')
     def default_dtypes(
         self,
         *,
-        device: Optional[device] = None,
+        device: Device | None = None,
     ) -> DefaultDataTypes:
         return {
-            "real floating": float64,
-            "complex floating": complex128,
-            "integral": int64,
-            "indexing": int64,
+            "real floating": dt.float64,
+            "complex floating": dt.complex128,
+            "integral": dt.int64,
+            "indexing": dt.int64,
         }
 
     @requires_api_version('2023.12')
     def dtypes(
         self,
         *,
-        device: Optional[device] = None,
-        kind: Optional[Union[str, Tuple[str, ...]]] = None,
+        device: Device | None = None,
+        kind: str | tuple[str, ...] | None = None,
path: array_api_strict/_info.py
user: login: jorenham
body: This is the kind of situation (pun intended) where a `Literal` can be helpful
created_at: 2025-03-21T19:53:11Z
author_association: MEMBER
diff_hunk: @@ -312,14 +317,14 @@ def slogdet(x: Array, /) -> SlogdetResult:
 
 # To workaround this, the below is the code from np.linalg.solve except
 # only calling solve1 in the exactly 1D case.
-def _solve(a, b):
+def _solve(a: np.ndarray, b: np.ndarray) -> np.ndarray:
     try:
-        from numpy.linalg._linalg import (
+        from numpy.linalg._linalg import (  # type: ignore[attr-defined]
path: array_api_strict/_linalg.py
user: login: jorenham
body: I opened https://github.com/numpy/numtype/issues/367 for this
created_at: 2025-03-21T20:00:53Z
author_association: MEMBER
diff_hunk: @@ -1,33 +1,31 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import partial
+from typing import Literal, NamedTuple
 
+import numpy as np
+import numpy.linalg
+
+from ._array_object import Array
+from ._data_type_functions import finfo
 from ._dtypes import (
+    DType,
     _floating_dtypes,
     _numeric_dtypes,
-    float32,
     complex64,
     complex128,
 )
-from ._data_type_functions import finfo
-from ._manipulation_functions import reshape
 from ._elementwise_functions import conj
-from ._array_object import Array
-from ._flags import requires_extension, get_array_api_strict_flags
+from ._flags import get_array_api_strict_flags, requires_extension
+from ._manipulation_functions import reshape
+from ._statistical_functions import _np_dtype_sumprod
 
 try:
-    from numpy._core.numeric import normalize_axis_tuple
+    from numpy._core.numeric import normalize_axis_tuple  # type: ignore[attr-defined]
path: array_api_strict/_linalg.py
user: login: jorenham
body: I opened https://github.com/numpy/numtype/issues/368 for this
created_at: 2025-03-21T20:03:46Z
author_association: MEMBER
diff_hunk: @@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
path: array_api_strict/_sorting_functions.py
user: login: jorenham
body: FYI, `numpy.sort` only considers the first letter of `kind` case-insensitively, so `"Q"` or `"QuantumSort300"` would also work. 

But since that can't be typed by something narrower than `str`, it currently only allows `{"Q", "quick", "quicksort"}` for quicksort and `{"S", "stable", "stablesort"}` for stable<sub>sort</sub>.

So for the sake of consistency, symmetry, and that vein on my forehead that's about to burst after seeing this, you might also want to consider something like

```suggestion
    kind: Literal["stablesort", "quicksort"] = "stablesort" if stable else "quicksort"
```

*entweder*

```suggestion
    kind: Literal["stable", "quick"] = "stable" if stable else "quick"
```

*soit*

```suggestion
    kind: Literal["S", "Q"] = "S" if stable else "Q"
```
created_at: 2025-03-21T20:15:33Z
author_association: MEMBER
diff_hunk: @@ -47,8 +50,7 @@ def sort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in sort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
-    res = np.sort(x._array, axis=axis, kind=kind)
+    res = np.sort(x._array, axis=axis, kind="stable" if stable else "quicksort")
path: array_api_strict/_sorting_functions.py
user: login: jorenham
body: me repito
created_at: 2025-03-21T20:16:43Z
author_association: MEMBER
diff_hunk: @@ -8,82 +8,66 @@
 
 from __future__ import annotations
 
-__all__ = [
-    "Array",
-    "Device",
-    "Dtype",
-    "SupportsDLPack",
-    "SupportsBufferProtocol",
-    "PyCapsule",
-]
-
 import sys
+from typing import Any, Protocol, TypedDict, TypeVar
 
-from typing import (
-    Any,
-    TypedDict,
-    TypeVar,
-    Protocol,
-)
-
-from ._array_object import Array, _device
-from ._dtypes import _DType
-from ._info import __array_namespace_info__
+from ._dtypes import DType
 
 _T_co = TypeVar("_T_co", covariant=True)
 
+
 class NestedSequence(Protocol[_T_co]):
     def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
     def __len__(self, /) -> int: ...
 
 
-Device = _device
-
-Dtype = _DType
-
-Info = __array_namespace_info__
-
 if sys.version_info >= (3, 12):
     from collections.abc import Buffer as SupportsBufferProtocol
 else:
     SupportsBufferProtocol = Any
 
 PyCapsule = Any
 
+
 class SupportsDLPack(Protocol):
     def __dlpack__(self, /, *, stream: None = ...) -> PyCapsule: ...
 
+
 Capabilities = TypedDict(
-    "Capabilities", {"boolean indexing": bool, "data-dependent shapes": bool,
-                     "max dimensions": int}
+    "Capabilities",
+    {
+        "boolean indexing": bool,
+        "data-dependent shapes": bool,
+        "max dimensions": int,
+    },
 )
 
 DefaultDataTypes = TypedDict(
     "DefaultDataTypes",
     {
-        "real floating": Dtype,
-        "complex floating": Dtype,
-        "integral": Dtype,
-        "indexing": Dtype,
+        "real floating": DType,
+        "complex floating": DType,
+        "integral": DType,
+        "indexing": DType,
     },
 )
 
 DataTypes = TypedDict(
     "DataTypes",
     {
-        "bool": Dtype,
-        "float32": Dtype,
-        "float64": Dtype,
-        "complex64": Dtype,
-        "complex128": Dtype,
-        "int8": Dtype,
-        "int16": Dtype,
-        "int32": Dtype,
-        "int64": Dtype,
-        "uint8": Dtype,
-        "uint16": Dtype,
-        "uint32": Dtype,
-        "uint64": Dtype,
+        "bool": DType,
+        "float32": DType,
+        "float64": DType,
+        "complex64": DType,
+        "complex128": DType,
+        "int8": DType,
+        "int16": DType,
+        "int32": DType,
+        "int64": DType,
+        "uint8": DType,
+        "uint16": DType,
+        "uint32": DType,
+        "uint64": DType,
path: array_api_strict/_typing.py
user: login: jorenham
body: if there are no spaces, the general recommendation is to use the class-based syntax instead

```py
class DataTypes(TypedDict, total=false):
	bool: DType
	float32: DType
	float64: DType
	complex64: DType
	complex128: DType
	int8: DType
	int16: DType
	int32: DType
	int64: DType
	uint8: DType
	uint16: DType
	uint32: DType
	uint64: DType
```

I think there's something in `ruff` that would complain about it, as well 🤷🏻 
created_at: 2025-03-21T20:21:29Z
author_association: MEMBER
diff_hunk: @@ -57,7 +56,7 @@ def diff(
     # currently specified.
 
     # NumPy does not support prepend=None or append=None
-    kwargs = dict(axis=axis, n=n)
+    kwargs: dict[str, Any] = dict(axis=axis, n=n)
path: array_api_strict/_utility_functions.py
user: login: jorenham
body: The cool kids are all constructing their dicts as `{...}` instead of `dict(...)` nowadays, so how about:

```suggestion
    kwargs: dict[str, int | npt.NDArray[Any]] = {"axis": axis, "n", n}
```

I sneakily narrowed the dict type as well, which should get you couple of extra swag-creds
created_at: 2025-03-21T20:25:37Z
author_association: MEMBER
diff_hunk: @@ -1,11 +1,9 @@
-from typing import Callable
-
 import pytest
 
 import array_api_strict as xp
 
 
-def p(func: Callable, *args, **kwargs):
+def p(func, *args, **kwargs):
path: array_api_strict/tests/test_validation.py
user: login: jorenham
body: Something like 

```suggestion
def p(func: Callable[..., object], *args: object, **kwargs: object) -> pytest.param:
```

seems just as easy to me 🤷🏻.

But that being said, I'm the type that would even stub his own toe.
created_at: 2025-03-21T20:30:17Z
author_association: MEMBER
diff_hunk: @@ -31,3 +31,16 @@ Repository = "https://github.com/data-apis/array-api-strict"
 [tool.setuptools_scm]
 version_file = "array_api_strict/_version.py"
 
+[tool.mypy]
+disallow_incomplete_defs = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+show_error_codes = true
+warn_redundant_casts = true
+warn_unused_ignores = true
+warn_unreachable = true
path: pyproject.toml
user: login: jorenham
body: I would also recommend `strict_bytes = true`, which disables the "promotion" of `memoryview` and `bytesarray` to `bytes` (similar to how `int` gets promoted to `float` and `float` to `complex`). This will be enabled by default in mypy 2.0.

It might also be a good idea to set `local_partial_types = true`, which will also be the default in mypy 2.0.

Setting `strict = true` is also a good idea (and it isn't as strict as the name suggests it is).
created_at: 2025-03-21T20:37:11Z
author_association: MEMBER
diff_hunk: @@ -1,18 +1,24 @@
-"""Private helper routines.
-"""
+"""Private helper routines."""
 
-from ._flags import get_array_api_strict_flags
+from __future__ import annotations
+
+from ._array_object import Array
 from ._dtypes import _dtype_categories
+from ._flags import get_array_api_strict_flags
 
 _py_scalars = (bool, int, float, complex)
 
 
-def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
-
+def _maybe_normalize_py_scalars(
+    x1: Array | complex,
+    x2: Array | complex,
+    dtype_category: str,
+    func_name: str,
+) -> tuple[Array, Array]:
     flags = get_array_api_strict_flags()
     if flags["api_version"] < "2024.12":
         # scalars will fail at the call site
-        return x1, x2
+        return x1, x2  # type: ignore[return-value]
path: array_api_strict/_helpers.py
user: login: crusaderky
body: cast would be wrong. We're deliberately returning a potentially wrong value with the intent to make the caller crash. cast is for when you know the variable is of a certain type but the linter doesn't know yet.
created_at: 2025-03-21T21:20:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -36,104 +33,104 @@ def capabilities(self) -> Capabilities:
         return res
 
     @requires_api_version('2023.12')
-    def default_device(self) -> device:
+    def default_device(self) -> Device:
         return CPU_DEVICE
 
     @requires_api_version('2023.12')
     def default_dtypes(
         self,
         *,
-        device: Optional[device] = None,
+        device: Device | None = None,
     ) -> DefaultDataTypes:
         return {
-            "real floating": float64,
-            "complex floating": complex128,
-            "integral": int64,
-            "indexing": int64,
+            "real floating": dt.float64,
+            "complex floating": dt.complex128,
+            "integral": dt.int64,
+            "indexing": dt.int64,
         }
 
     @requires_api_version('2023.12')
     def dtypes(
         self,
         *,
-        device: Optional[device] = None,
-        kind: Optional[Union[str, Tuple[str, ...]]] = None,
+        device: Device | None = None,
+        kind: str | tuple[str, ...] | None = None,
path: array_api_strict/_info.py
user: login: crusaderky
body: I'd rather not deviate from the API standard.
created_at: 2025-03-21T21:24:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,33 +1,31 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import partial
+from typing import Literal, NamedTuple
 
+import numpy as np
+import numpy.linalg
+
+from ._array_object import Array
+from ._data_type_functions import finfo
 from ._dtypes import (
+    DType,
     _floating_dtypes,
     _numeric_dtypes,
-    float32,
     complex64,
     complex128,
 )
-from ._data_type_functions import finfo
-from ._manipulation_functions import reshape
 from ._elementwise_functions import conj
-from ._array_object import Array
-from ._flags import requires_extension, get_array_api_strict_flags
+from ._flags import get_array_api_strict_flags, requires_extension
+from ._manipulation_functions import reshape
+from ._statistical_functions import _np_dtype_sumprod
 
 try:
-    from numpy._core.numeric import normalize_axis_tuple
+    from numpy._core.numeric import normalize_axis_tuple  # type: ignore[attr-defined]
path: array_api_strict/_linalg.py
user: login: crusaderky
body: Aren't they going to answer that you're not meant to import private API?
created_at: 2025-03-21T21:24:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
path: array_api_strict/_sorting_functions.py
user: login: crusaderky
body: The numpy API specification says 

> kind{‘quicksort’, ‘mergesort’, ‘heapsort’, ‘stable’}, optional
created_at: 2025-03-21T21:26:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,11 +1,9 @@
-from typing import Callable
-
 import pytest
 
 import array_api_strict as xp
 
 
-def p(func: Callable, *args, **kwargs):
+def p(func, *args, **kwargs):
path: array_api_strict/tests/test_validation.py
user: login: crusaderky
body: Hard no from me. This is a test module.
created_at: 2025-03-21T21:28:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -31,3 +31,16 @@ Repository = "https://github.com/data-apis/array-api-strict"
 [tool.setuptools_scm]
 version_file = "array_api_strict/_version.py"
 
+[tool.mypy]
+disallow_incomplete_defs = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+show_error_codes = true
+warn_redundant_casts = true
+warn_unused_ignores = true
+warn_unreachable = true
path: pyproject.toml
user: login: crusaderky
body: I'd rather treat these settings as a good starting point (I copy-pasted them from my template project) and leave tweaking them to a follow-up PR. There are many unfixed issues with them already.
created_at: 2025-03-21T21:30:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -31,3 +31,16 @@ Repository = "https://github.com/data-apis/array-api-strict"
 [tool.setuptools_scm]
 version_file = "array_api_strict/_version.py"
 
+[tool.mypy]
+disallow_incomplete_defs = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+show_error_codes = true
+warn_redundant_casts = true
+warn_unused_ignores = true
+warn_unreachable = true
path: pyproject.toml
user: login: jorenham
body: I don't expect that enabling `strict_bytes` and `local_partial_types` will lead to different outcomes in this project.
created_at: 2025-03-22T16:32:51Z
author_association: MEMBER
diff_hunk: @@ -1,11 +1,9 @@
-from typing import Callable
-
 import pytest
 
 import array_api_strict as xp
 
 
-def p(func: Callable, *args, **kwargs):
+def p(func, *args, **kwargs):
path: array_api_strict/tests/test_validation.py
user: login: jorenham
body: > This is a test module.

Yea I noticed that.

> Hard no from me.

Why?
created_at: 2025-03-22T16:37:26Z
author_association: MEMBER
diff_hunk: @@ -1,18 +1,24 @@
-"""Private helper routines.
-"""
+"""Private helper routines."""
 
-from ._flags import get_array_api_strict_flags
+from __future__ import annotations
+
+from ._array_object import Array
 from ._dtypes import _dtype_categories
+from ._flags import get_array_api_strict_flags
 
 _py_scalars = (bool, int, float, complex)
 
 
-def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
-
+def _maybe_normalize_py_scalars(
+    x1: Array | complex,
+    x2: Array | complex,
+    dtype_category: str,
+    func_name: str,
+) -> tuple[Array, Array]:
     flags = get_array_api_strict_flags()
     if flags["api_version"] < "2024.12":
         # scalars will fail at the call site
-        return x1, x2
+        return x1, x2  # type: ignore[return-value]
path: array_api_strict/_helpers.py
user: login: jorenham
body: > cast is for when you know the variable is of a certain type but the linter doesn't know yet

For array-api<2024.15 the implicit assumption is that the input is not scalar. The type-checker doesn't know that. So doesn't that justify it?
created_at: 2025-03-22T16:40:46Z
author_association: MEMBER
diff_hunk: @@ -36,104 +33,104 @@ def capabilities(self) -> Capabilities:
         return res
 
     @requires_api_version('2023.12')
-    def default_device(self) -> device:
+    def default_device(self) -> Device:
         return CPU_DEVICE
 
     @requires_api_version('2023.12')
     def default_dtypes(
         self,
         *,
-        device: Optional[device] = None,
+        device: Device | None = None,
     ) -> DefaultDataTypes:
         return {
-            "real floating": float64,
-            "complex floating": complex128,
-            "integral": int64,
-            "indexing": int64,
+            "real floating": dt.float64,
+            "complex floating": dt.complex128,
+            "integral": dt.int64,
+            "indexing": dt.int64,
         }
 
     @requires_api_version('2023.12')
     def dtypes(
         self,
         *,
-        device: Optional[device] = None,
-        kind: Optional[Union[str, Tuple[str, ...]]] = None,
+        device: Device | None = None,
+        kind: str | tuple[str, ...] | None = None,
path: array_api_strict/_info.py
user: login: jorenham
body: The API explicitly states that if `kind` is a string, is must be one of the listed string literals. In that sense, using a literal would make it deviate *less* from the API standard.
https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.dtypes.html
created_at: 2025-03-22T16:44:39Z
author_association: MEMBER
diff_hunk: @@ -1,33 +1,31 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import partial
+from typing import Literal, NamedTuple
 
+import numpy as np
+import numpy.linalg
+
+from ._array_object import Array
+from ._data_type_functions import finfo
 from ._dtypes import (
+    DType,
     _floating_dtypes,
     _numeric_dtypes,
-    float32,
     complex64,
     complex128,
 )
-from ._data_type_functions import finfo
-from ._manipulation_functions import reshape
 from ._elementwise_functions import conj
-from ._array_object import Array
-from ._flags import requires_extension, get_array_api_strict_flags
+from ._flags import get_array_api_strict_flags, requires_extension
+from ._manipulation_functions import reshape
+from ._statistical_functions import _np_dtype_sumprod
 
 try:
-    from numpy._core.numeric import normalize_axis_tuple
+    from numpy._core.numeric import normalize_axis_tuple  # type: ignore[attr-defined]
path: array_api_strict/_linalg.py
user: login: jorenham
body: "They" in this context refers to a guy that's very sympathetic to the array API projects <sub>who's also very handsome and has a great sense of humor</sub>.

Coincidentally, "they" just merged a PR that fixes this:

https://github.com/numpy/numtype/pull/371
created_at: 2025-03-22T16:49:13Z
author_association: MEMBER
diff_hunk: @@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
path: array_api_strict/_sorting_functions.py
user: login: jorenham
body: Yes, but that's incomplete:

```pycon
>>> import numpy as np
>>> np.sort([1,3,2], kind="S")
array([1, 2, 3])
>>> np.sort([1,3,2], kind="Q")
array([1, 2, 3])
>>> np.sort([1,3,2], kind="stablesort")
array([1, 2, 3])
>>> np.sort([1,3,2], kind="quick'")
array([1, 2, 3])
```
created_at: 2025-03-22T16:50:42Z
author_association: MEMBER
diff_hunk: @@ -16,62 +16,67 @@
 from __future__ import annotations
 
 import operator
+from collections.abc import Iterator
 from enum import IntEnum
+from types import ModuleType
+from typing import Any, Literal, SupportsIndex
 
-from ._creation_functions import asarray
+import numpy as np
+import numpy.typing as npt
+
+from ._creation_functions import _default, _Default, asarray
 from ._dtypes import (
-    _DType,
+    DType,
     _all_dtypes,
     _boolean_dtypes,
+    _complex_floating_dtypes,
+    _dtype_categories,
+    _floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
-    _floating_dtypes,
-    _real_floating_dtypes,
-    _complex_floating_dtypes,
     _numeric_dtypes,
-    _result_type,
-    _dtype_categories,
+    _real_floating_dtypes,
     _real_to_complex_map,
+    _result_type,
 )
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
+from ._typing import PyCapsule
 
-from typing import TYPE_CHECKING, SupportsIndex
-import types
+try:
+    from types import EllipsisType  # Python >=3.10
+except ImportError:
+    EllipsisType = type(Ellipsis)
 
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union, Any
-    from ._typing import PyCapsule, Dtype
-    import numpy.typing as npt
-
-import numpy as np
 
 class Device:
-    def __init__(self, device="CPU_DEVICE"):
+    _device: str
path: array_api_strict/_array_object.py
user: login: crusaderky
body: I see no reason to prevent another library from subclassing this
created_at: 2025-03-26T14:06:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,18 +1,24 @@
-"""Private helper routines.
-"""
+"""Private helper routines."""
 
-from ._flags import get_array_api_strict_flags
+from __future__ import annotations
+
+from ._array_object import Array
 from ._dtypes import _dtype_categories
+from ._flags import get_array_api_strict_flags
 
 _py_scalars = (bool, int, float, complex)
 
 
-def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
-
+def _maybe_normalize_py_scalars(
+    x1: Array | complex,
+    x2: Array | complex,
+    dtype_category: str,
+    func_name: str,
+) -> tuple[Array, Array]:
     flags = get_array_api_strict_flags()
     if flags["api_version"] < "2024.12":
         # scalars will fail at the call site
-        return x1, x2
+        return x1, x2  # type: ignore[return-value]
path: array_api_strict/_helpers.py
user: login: crusaderky
body: cast is still wrong unless I also update the return type of the function.
Which I'd rather not do, as I'd rather not complicate typing annotations for the benefit of obsolete versions of the API.

created_at: 2025-03-26T14:19:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -36,104 +33,104 @@ def capabilities(self) -> Capabilities:
         return res
 
     @requires_api_version('2023.12')
-    def default_device(self) -> device:
+    def default_device(self) -> Device:
         return CPU_DEVICE
 
     @requires_api_version('2023.12')
     def default_dtypes(
         self,
         *,
-        device: Optional[device] = None,
+        device: Device | None = None,
     ) -> DefaultDataTypes:
         return {
-            "real floating": float64,
-            "complex floating": complex128,
-            "integral": int64,
-            "indexing": int64,
+            "real floating": dt.float64,
+            "complex floating": dt.complex128,
+            "integral": dt.int64,
+            "indexing": dt.int64,
         }
 
     @requires_api_version('2023.12')
     def dtypes(
         self,
         *,
-        device: Optional[device] = None,
-        kind: Optional[Union[str, Tuple[str, ...]]] = None,
+        device: Device | None = None,
+        kind: str | tuple[str, ...] | None = None,
path: array_api_strict/_info.py
user: login: crusaderky
body: This must match `isdtype`, which is a _very_ widely used function. This is a nuanced discussion that I'd rather leave to a follow-up PR.
created_at: 2025-03-26T14:22:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
path: array_api_strict/_sorting_functions.py
user: login: crusaderky
body: We should not care if a specific backend offers more than what the Array API requires.
created_at: 2025-03-26T14:23:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,11 +1,9 @@
-from typing import Callable
-
 import pytest
 
 import array_api_strict as xp
 
 
-def p(func: Callable, *args, **kwargs):
+def p(func, *args, **kwargs):
path: array_api_strict/tests/test_validation.py
user: login: crusaderky
body: type annotations in test modules unnecessarily increase attrition and code verbosity without giving any material benefit to code maintainability or robustness. A pathological example of that is array-api-extra.
created_at: 2025-03-26T14:27:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,18 +1,24 @@
-"""Private helper routines.
-"""
+"""Private helper routines."""
 
-from ._flags import get_array_api_strict_flags
+from __future__ import annotations
+
+from ._array_object import Array
 from ._dtypes import _dtype_categories
+from ._flags import get_array_api_strict_flags
 
 _py_scalars = (bool, int, float, complex)
 
 
-def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
-
+def _maybe_normalize_py_scalars(
+    x1: Array | complex,
+    x2: Array | complex,
+    dtype_category: str,
+    func_name: str,
+) -> tuple[Array, Array]:
     flags = get_array_api_strict_flags()
     if flags["api_version"] < "2024.12":
         # scalars will fail at the call site
-        return x1, x2
+        return x1, x2  # type: ignore[return-value]
path: array_api_strict/_helpers.py
user: login: jorenham
body: fair enough
created_at: 2025-03-26T18:48:56Z
author_association: MEMBER
diff_hunk: @@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
path: array_api_strict/_sorting_functions.py
user: login: jorenham
body: I'd agree with you if it would've been consistent on its own, which it isn't in this case.
created_at: 2025-03-26T18:50:15Z
author_association: MEMBER
diff_hunk: @@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
path: array_api_strict/_sorting_functions.py
user: login: crusaderky
body: I can change to `kind: Literal['quicksort', 'mergesort', 'heapsort', 'stable']` if you prefer
created_at: 2025-03-26T18:58:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,11 +1,9 @@
-from typing import Callable
-
 import pytest
 
 import array_api_strict as xp
 
 
-def p(func: Callable, *args, **kwargs):
+def p(func, *args, **kwargs):
path: array_api_strict/tests/test_validation.py
user: login: jorenham
body: I don't really agree with that. 

Tests have to be maintained as well, and annotations make it easier for the maintainer to see what the *intent* of the tests is. 

In case of test failures, it can be very helpful to see the annotations printed in the error output, because in many cases it is helpful to know the *intent* of the tests when you are figuring out when it fails.

So I guess I view type annotations in tests as a concise yet precise form of documentation, specifically meant to make our own (i.e. the maintainers) lives a bit easier.

Either way, do whatever floats your boat. I just know that it has helped in my projects, so I think it can help here as well.
created_at: 2025-03-26T18:59:46Z
author_association: MEMBER
diff_hunk: @@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
path: array_api_strict/_sorting_functions.py
user: login: jorenham
body: if you remove the `*sort` bits, then your got yourself a deal 😉 
created_at: 2025-03-26T19:19:32Z
author_association: MEMBER
diff_hunk: @@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
path: array_api_strict/_sorting_functions.py
user: login: crusaderky
body: I find it would be bonkers to prefer an undocumented variant, which potentially may become deprecated in the future, to the official one.

The [official API documentation](https://numpy.org/doc/2.1/reference/generated/numpy.argsort.html) states:

> kind{‘quicksort’, ‘mergesort’, ‘heapsort’, ‘stable’}, optional
> Sorting algorithm. The default is ‘quicksort’. Note that both ‘stable’ and ‘mergesort’ use timsort under the covers and, in general, the actual implementation will vary with data type. The ‘mergesort’ option is retained for backwards compatibility.

There is no mention whatsoever of the other variants accepted in [_SortKind](https://github.com/numpy/numpy/blob/9389862162bbd46b5324402d6a08a27bc18ddb7d/numpy/__init__.pyi#L993-L998), which just happen to be accepted I assume out of backwards compatibility concerns.


created_at: 2025-03-27T13:28:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
path: array_api_strict/_sorting_functions.py
user: login: jorenham
body: > I find it would be bonkers to prefer an undocumented variant, which potentially may become deprecated in the future, to the official one.

Well, it's no secret that I'm bonkers, but I'm not sure how that's relevant here 🤷🏻 

And yes, it might become deprecated, just like anything that's documented might.

> The [official API documentation](https://numpy.org/doc/2.1/reference/generated/numpy.argsort.html) states:

The [*other* official API documentation](https://github.com/numpy/numpy/blob/main/numpy/__init__.pyi#L991-L998), often referred to as the stubs, states that it *is* supported: 

```pyi
# in practice, only the first case-insensitive character is considered (so e.g.
# "QuantumSort3000" will be interpreted as quicksort).
_SortKind: TypeAlias = L[
    "Q", "quick", "quicksort",
    "M", "merge", "mergesort",
    "H", "heap", "heapsort",
    "S", "stable", "stablesort",
]
```

But it's also fine by me if you do keep it like this. Because as I've said before, it's  just a suggestion. No need to worry about that vein on my forehead ;)
created_at: 2025-03-27T13:54:36Z
author_association: MEMBER
diff_hunk: @@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
path: array_api_strict/_sorting_functions.py
user: login: crusaderky
body: Stubs are *not* documentation.
Stubs are an optional clutch to detach type annotations from the code. They should otherwise be treated as part of the implementation.

_The_ (one and only) API documentation does not mention those extra possible values. Which means they are an implementation detail.
created_at: 2025-03-27T17:35:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
path: array_api_strict/_sorting_functions.py
user: login: jorenham
body: I disagree (but I'm probably the most biased person you could find). And since there's no "agree to disagree" button, marking it as "resolved" is fine by me

created_at: 2025-03-28T14:44:00Z
author_association: MEMBER
diff_hunk: @@ -680,14 +701,14 @@ def __ge__(self: Array, other: Union[int, float, Array], /) -> Array:
         return self.__class__._new(res, device=self.device)
 
     def __getitem__(
-        self: Array,
-        key: Union[
-            int,
-            slice,
-            ellipsis,  # noqa: F821
-            Tuple[Union[int, slice, ellipsis, None], ...],  # noqa: F821
-            Array,
-        ],
+        self,
+        key: (
+            int
+            | slice
+            | EllipsisType
+            | tuple[int | slice | EllipsisType | None, ...]
+            | Array
+        ),
path: array_api_strict/_array_object.py
user: login: lucascolley
body: > to express it in full we'll need way way more (overloads?)

forgive my ignorance, why can't we just add `Array` into the `tuple` to fix this problem?
created_at: 2025-03-29T21:52:25Z
author_association: MEMBER
diff_hunk: @@ -680,14 +701,14 @@ def __ge__(self: Array, other: Union[int, float, Array], /) -> Array:
         return self.__class__._new(res, device=self.device)
 
     def __getitem__(
-        self: Array,
-        key: Union[
-            int,
-            slice,
-            ellipsis,  # noqa: F821
-            Tuple[Union[int, slice, ellipsis, None], ...],  # noqa: F821
-            Array,
-        ],
+        self,
+        key: (
+            int
+            | slice
+            | EllipsisType
+            | tuple[int | slice | EllipsisType | None, ...]
+            | Array
+        ),
path: array_api_strict/_array_object.py
user: login: crusaderky
body: Fixed
created_at: 2025-03-31T13:22:04Z
author_association: CONTRIBUTOR

PR comments:
user: login: ev-br
created_at: 2025-03-21T15:43:36Z
body: Can ignore python 3.9, cf https://github.com/data-apis/array-api-strict/pull/129
author_association: MEMBER
user: login: crusaderky
created_at: 2025-03-21T16:56:12Z
body: > Can ignore python 3.9, cf #129

Unless you plan to merge #129 now, I can't ignore 3.9 as it would make _all_ CI to fail.
author_association: CONTRIBUTOR
user: login: crusaderky
created_at: 2025-03-21T17:00:38Z
body: CC @jorenham  @NeilGirdhar 
author_association: CONTRIBUTOR
user: login: jorenham
created_at: 2025-03-21T20:45:07Z
body: I took the liberty of creating a label for static typing stuff <sub>and painted it black <sub>because that's how I prefer my coffee <sub>(and my metal)</sub></sub></sub>
author_association: MEMBER
user: login: crusaderky
created_at: 2025-03-26T16:21:10Z
body: > Let me record my -1 on `bool | int |float` until [data-apis/array-api-compat#257 (comment)](https://github.com/data-apis/array-api-compat/pull/257#issuecomment-2746365323) is done.

Can we please find a minimum common ground that allows us to merge this PR?
I've just changed all scalar types to match exactly the Array API docs. That means explicit `bool | int | float | complex`, `bool | int`, and `int | float | complex`.
author_association: CONTRIBUTOR
user: login: ev-br
created_at: 2025-03-26T16:40:00Z
body: To repeat, my only strong feelings are on 
- consistency of annotations and signatures with the Array API docs and stubs, and
- `int | bool | float`,
in this order.

With these two items being there, I'm declaring my expertise on typing to run out, the typing itself being considered mostly harmless; I'm thus dismissing my request for changes and bow out. As long as my boundary conditions are not exceeded, I'm happy with whatever y'all converge on. 
author_association: MEMBER
user: login: crusaderky
created_at: 2025-03-26T16:46:01Z
body: This is ready for final review.
author_association: CONTRIBUTOR
user: login: crusaderky
created_at: 2025-03-27T17:36:38Z
body: I believe there should not be any oustanding comments?
author_association: CONTRIBUTOR
user: login: lucascolley
created_at: 2025-03-29T22:15:34Z
body: ### Trying to push things forward hat on:

It sounds like both @crusaderky and I would be happy to keep `int | bool | float` in for now given @ev-br's concern, with a view to eventually removing them once we have array-api-typing up and running and readable/documented aliases, which make sense appearing both in the 'real' annotations and on the spec pages.

Would merging @crusaderky's 2 PRs disrupt your work at https://github.com/data-apis/array-api-compat/pull/288 @jorenham ?


> - consistency of annotations and signatures with the Array API docs and stubs, and
> - `int | bool | float`,
> in this order.

I haven't read the entirety of all of these threads, but are there any outstanding concerns about consistency *apart from* the `int | bool | float` situation?

---

### Opinion hat on:

For me, there is no significant consistency concern about `int | bool | float`, for reasons already stressed in this thread: these annotations are tool/developer facing, not user facing. I think it is fairly obvious to developers adopting the spec that the spec pages are the source of truth. Following SciPy's lead, we should distinguish between documented types (for docs readability) and the actual type annotations, keeping in mind that we can hope for these to converge somewhat with array-api-typing. In this sense some 'inconsistency' is actually desirable.

If the concern is that inconsistency between the spec saying `int | float` and an annotation in (say) array-api-strict of `float` will confuse array-api-strict developers, I think this is best remedied in the relevant contributor guides, rather than holding up the typing experts from making the much-desired improvements.
author_association: MEMBER
user: login: jorenham
created_at: 2025-03-31T22:26:13Z
body: > Would merging @crusaderky's 2 PRs disrupt your work at [data-apis/array-api-compat#288](https://github.com/data-apis/array-api-compat/pull/288) @jorenham ?

maybe a little bit; but I'll survive
author_association: MEMBER
user: login: lucascolley
created_at: 2025-03-31T22:32:11Z
body: okay, let's make some progress here then - thanks as always Joren, and Guido for the PR
author_association: MEMBER