array-api-strict 2.3__py3-none-any.whl → 2.4__py3-none-any.whl

array_api_strict/_linalg.py

@@ -1,33 +1,23 @@
-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 ._dtypes import (
-    _floating_dtypes,
-    _numeric_dtypes,
-    float32,
-    complex64,
-    complex128,
-)
+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 typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from ._typing import Literal, Optional, Sequence, Tuple, Union, Dtype
+    from numpy.core.numeric import normalize_axis_tuple  # type: ignore[no-redef]
 
-from typing import NamedTuple
-
-import numpy.linalg
-import numpy as np
 
 class EighResult(NamedTuple):
     eigenvalues: Array
@@ -175,7 +165,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 +182,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 +205,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 +217,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 +250,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 +265,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 +309,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 +379,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 +395,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 +409,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: float = 2,
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.linalg.norm <numpy.linalg.norm>`.
 
@@ -456,8 +454,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 +478,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)
 

array_api_strict/_linear_algebra_functions.py

@@ -4,19 +4,15 @@ them here with wrappers in linalg so that the wrappers can be disabled if the
 linalg extension is disabled in the flags.
 
 """
+from collections.abc import Sequence
 
-from __future__ import annotations
+import numpy as np
+import numpy.linalg
 
-from ._dtypes import _numeric_dtypes
 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 +34,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:

array_api_strict/_manipulation_functions.py

@@ -1,21 +1,15 @@
-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 +23,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 +39,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 +50,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 +63,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 +74,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,17 +105,16 @@ 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>`.
 
     See its docstring for more information.
     """
+    if not isinstance(shape, tuple):
+        raise TypeError(f"`shape` must be a tuple of ints; got {shape=} instead.")
 
     data = x._array
     if copy:
@@ -135,19 +131,23 @@ 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>`.
 
     See its docstring for more information.
     """
+    if not isinstance(shift, int | tuple):
+        raise ValueError(
+            f"`shift` can only be an int or a tuple, got {type(shift)=} instead."
+        )
     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 +161,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 +172,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 +190,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")
 

array_api_strict/_searching_functions.py

@@ -1,17 +1,15 @@
-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 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>`.
 
@@ -22,7 +20,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>`.
 
@@ -34,7 +32,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>`.
 
@@ -51,7 +49,7 @@ def count_nonzero(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     """
@@ -70,7 +68,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>`.
@@ -83,36 +81,27 @@ 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, /
-) -> Array:
+
+def where(condition: Array, x1: Array | complex, x2: Array | 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':
-        num_scalars = 0
-
-        if isinstance(x1, (bool, float, complex, int)):
-            x1 = Array._new(np.asarray(x1), device=condition.device)
-            num_scalars += 1
-
-        if isinstance(x2, (bool, float, complex, int)):
-            x2 = Array._new(np.asarray(x2), device=condition.device)
-            num_scalars += 1
+    if not isinstance(condition, Array):
+        raise TypeError(f"`condition` must be an Array; got {type(condition)}")
 
-        if num_scalars == 2:
-            raise ValueError("One of x1, x2 arguments must be an array.")
+    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)

array_api_strict/_set_functions.py

@@ -1,13 +1,10 @@
-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 +17,7 @@ import numpy as np
 # Note: The functions here return a namedtuple (np.unique() returns a normal
 # tuple).
 
+
 class UniqueAllResult(NamedTuple):
     values: Array
     indices: Array

array_api_strict/_sorting_functions.py

@@ -1,10 +1,10 @@
-from __future__ import annotations
+from typing import Literal
+
+import numpy as np
 
 from ._array_object import Array
 from ._dtypes import _real_numeric_dtypes
 
-import numpy as np
-
 
 # Note: the descending keyword argument is new in this function
 def argsort(
@@ -18,7 +18,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 +35,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 +48,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)