TYP: fix typing errors in common._linalg

Author: jorenham

array_api_compat/common/_linalg.py

@@ -1,23 +1,33 @@
 from __future__ import annotations
 
 import math
-from typing import Literal, NamedTuple, Optional, Tuple, Union
+from typing import Literal, NamedTuple, cast
 
 import numpy as np
+
 if np.__version__[0] == "2":
     from numpy.lib.array_utils import normalize_axis_tuple
 else:
     from numpy.core.numeric import normalize_axis_tuple
 
-from ._aliases import matmul, matrix_transpose, tensordot, vecdot, isdtype
 from .._internal import get_xp
-from ._typing import Array, Namespace
+from ._aliases import isdtype, matmul, matrix_transpose, tensordot, vecdot
+from ._typing import Array, DType, Namespace
+
 
 # These are in the main NumPy namespace but not in numpy.linalg
-def cross(x1: Array, x2: Array, /, xp: Namespace, *, axis: int = -1, **kwargs) -> Array:
+def cross(
+    x1: Array,
+    x2: Array,
+    /,
+    xp: Namespace,
+    *,
+    axis: int = -1,
+    **kwargs: object,
+) -> Array:
     return xp.cross(x1, x2, axis=axis, **kwargs)
 
-def outer(x1: Array, x2: Array, /, xp: Namespace, **kwargs) -> Array:
+def outer(x1: Array, x2: Array, /, xp: Namespace, **kwargs: object) -> Array:
     return xp.outer(x1, x2, **kwargs)
 
 class EighResult(NamedTuple):
@@ -39,46 +49,66 @@ class SVDResult(NamedTuple):
 
 # These functions are the same as their NumPy counterparts except they return
 # a namedtuple.
-def eigh(x: Array, /, xp: Namespace, **kwargs) -> EighResult:
+def eigh(x: Array, /, xp: Namespace, **kwargs: object) -> EighResult:
     return EighResult(*xp.linalg.eigh(x, **kwargs))
 
-def qr(x: Array, /, xp: Namespace, *, mode: Literal['reduced', 'complete'] = 'reduced',
-       **kwargs) -> QRResult:
+def qr(
+    x: Array,
+    /,
+    xp: Namespace,
+    *,
+    mode: Literal["reduced", "complete"] = "reduced",
+    **kwargs: object,
+) -> QRResult:
     return QRResult(*xp.linalg.qr(x, mode=mode, **kwargs))
 
-def slogdet(x: Array, /, xp: Namespace, **kwargs) -> SlogdetResult:
+def slogdet(x: Array, /, xp: Namespace, **kwargs: object) -> SlogdetResult:
     return SlogdetResult(*xp.linalg.slogdet(x, **kwargs))
 
 def svd(
-    x: Array, /, xp: Namespace, *, full_matrices: bool = True, **kwargs
+    x: Array,
+    /,
+    xp: Namespace,
+    *,
+    full_matrices: bool = True,
+    **kwargs: object,
 ) -> SVDResult:
     return SVDResult(*xp.linalg.svd(x, full_matrices=full_matrices, **kwargs))
 
 # These functions have additional keyword arguments
 
 # The upper keyword argument is new from NumPy
-def cholesky(x: Array, /, xp: Namespace, *, upper: bool = False, **kwargs) -> Array:
+def cholesky(
+    x: Array,
+    /,
+    xp: Namespace,
+    *,
+    upper: bool = False,
+    **kwargs: object,
+) -> Array:
     L = xp.linalg.cholesky(x, **kwargs)
     if upper:
         U = get_xp(xp)(matrix_transpose)(L)
         if get_xp(xp)(isdtype)(U.dtype, 'complex floating'):
-            U = xp.conj(U)
+            U = xp.conj(U)  # pyright: ignore[reportConstantRedefinition]
         return U
     return L
 
 # The rtol keyword argument of matrix_rank() and pinv() is new from NumPy.
 # Note that it has a different semantic meaning from tol and rcond.
-def matrix_rank(x: Array,
-                /,
-                xp: Namespace,
-                *,
-                rtol: Optional[Union[float, Array]] = None,
-                **kwargs) -> Array:
+def matrix_rank(
+    x: Array,
+    /,
+    xp: Namespace,
+    *,
+    rtol: float | Array | None = None,
+    **kwargs: object,
+) -> Array:
     # this is different from xp.linalg.matrix_rank, which supports 1
     # dimensional arrays.
     if x.ndim < 2:
         raise xp.linalg.LinAlgError("1-dimensional array given. Array must be at least two-dimensional")
-    S = get_xp(xp)(svdvals)(x, **kwargs)
+    S: Array = get_xp(xp)(svdvals)(x, **kwargs)
     if rtol is None:
         tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * xp.finfo(S.dtype).eps
     else:
@@ -88,7 +118,12 @@ def matrix_rank(x: Array,
     return xp.count_nonzero(S > tol, axis=-1)
 
 def pinv(
-    x: Array, /, xp: Namespace, *, rtol: Optional[Union[float, Array]] = None, **kwargs
+    x: Array,
+    /,
+    xp: Namespace,
+    *,
+    rtol: float | Array | None = None,
+    **kwargs: object,
 ) -> Array:
     # this is different from xp.linalg.pinv, which does not multiply the
     # default tolerance by max(M, N).
@@ -104,23 +139,23 @@ def matrix_norm(
     xp: Namespace,
     *,
     keepdims: bool = False,
-    ord: Optional[Union[int, float, Literal['fro', 'nuc']]] = 'fro',
+    ord: float | Literal["fro", "nuc"] | None = "fro",
 ) -> Array:
     return xp.linalg.norm(x, axis=(-2, -1), keepdims=keepdims, ord=ord)
 
 # svdvals is not in NumPy (but it is in SciPy). It is equivalent to
 # xp.linalg.svd(compute_uv=False).
-def svdvals(x: Array, /, xp: Namespace) -> Union[Array, Tuple[Array, ...]]:
+def svdvals(x: Array, /, xp: Namespace) -> Array | tuple[Array, ...]:
     return xp.linalg.svd(x, compute_uv=False)
 
 def vector_norm(
     x: Array,
     /,
     xp: Namespace,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
-    ord: Optional[Union[int, float]] = 2,
+    ord: float = 2,
 ) -> Array:
     # xp.linalg.norm tries to do a matrix norm whenever axis is a 2-tuple or
     # when axis=None and the input is 2-D, so to force a vector norm, we make
@@ -133,7 +168,10 @@ def vector_norm(
     elif isinstance(axis, tuple):
         # Note: The axis argument supports any number of axes, whereas
         # xp.linalg.norm() only supports a single axis for vector norm.
-        normalized_axis = normalize_axis_tuple(axis, x.ndim)
+        normalized_axis = cast(
+            "tuple[int, ...]",
+            normalize_axis_tuple(axis, x.ndim),  # pyright: ignore[reportCallIssue]
+        )
         rest = tuple(i for i in range(x.ndim) if i not in normalized_axis)
         newshape = axis + rest
         _x = xp.transpose(x, newshape).reshape(
@@ -149,7 +187,13 @@ def vector_norm(
         # We can't reuse xp.linalg.norm(keepdims) because of the reshape hacks
         # above to avoid matrix norm logic.
         shape = list(x.shape)
-        _axis = normalize_axis_tuple(range(x.ndim) if axis is None else axis, x.ndim)
+        _axis = cast(
+            "tuple[int, ...]",
+            normalize_axis_tuple(  # pyright: ignore[reportCallIssue]
+                range(x.ndim) if axis is None else axis,
+                x.ndim,
+            ),
+        )
         for i in _axis:
             shape[i] = 1
         res = xp.reshape(res, tuple(shape))
@@ -159,11 +203,17 @@ def vector_norm(
 # xp.diagonal and xp.trace operate on the first two axes whereas these
 # operates on the last two
 
-def diagonal(x: Array, /, xp: Namespace, *, offset: int = 0, **kwargs) -> Array:
+def diagonal(x: Array, /, xp: Namespace, *, offset: int = 0, **kwargs: object) -> Array:
     return xp.diagonal(x, offset=offset, axis1=-2, axis2=-1, **kwargs)
 
 def trace(
-    x: Array, /, xp: Namespace, *, offset: int = 0, dtype=None, **kwargs
+    x: Array,
+    /,
+    xp: Namespace,
+    *,
+    offset: int = 0,
+    dtype: DType | None = None,
+    **kwargs: object,
 ) -> Array:
     return xp.asarray(
         xp.trace(x, offset=offset, dtype=dtype, axis1=-2, axis2=-1, **kwargs)