Add matrix_transpose, svdvals, and vecdot

Author: asmeurer

numpy_array_api_compat/linalg.py

@@ -11,6 +11,33 @@
 def matrix_norm(x: ndarray, /, *, keepdims: bool = False, ord: Optional[Union[int, float, Literal['fro', 'nuc']]] = 'fro') -> ndarray:
     return np.linalg.norm(x, axis=(-2, -1), keepdims=keepdims, ord=ord)
 
+# this function is new in the array API spec. Unlike transpose, it only
+# transposes the last two axes.
+def matrix_transpose(x: ndarray, /) -> ndarray:
+    if x.ndim < 2:
+        raise ValueError("x must be at least 2-dimensional for matrix_transpose")
+    return np.swapaxes(x, -1, -2)
+
+# svdvals is not in NumPy (but it is in SciPy). It is equivalent to
+# np.linalg.svd(compute_uv=False).
+def svdvals(x: ndarray, /) -> Union[ndarray, Tuple[ndarray, ...]]:
+    return np.linalg.svd(x, compute_uv=False)
+
+# vecdot is not in NumPy
+def vecdot(x1: ndarray, x2: ndarray, /, *, axis: int = -1) -> ndarray:
+    ndim = max(x1.ndim, x2.ndim)
+    x1_shape = (1,)*(ndim - x1.ndim) + tuple(x1.shape)
+    x2_shape = (1,)*(ndim - x2.ndim) + tuple(x2.shape)
+    if x1_shape[axis] != x2_shape[axis]:
+        raise ValueError("x1 and x2 must have the same size along the given axis")
+
+    x1_, x2_ = np.broadcast_arrays(x1, x2)
+    x1_ = np.moveaxis(x1_, axis, -1)
+    x2_ = np.moveaxis(x2_, axis, -1)
+
+    res = x1_[..., None, :] @ x2_[..., None]
+    return res[..., 0, 0]
+
 def vector_norm(x: ndarray, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False, ord: Optional[Union[int, float]] = 2) -> ndarray:
     # np.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
@@ -52,5 +79,5 @@ def vector_norm(x: ndarray, /, *, axis: Optional[Union[int, Tuple[int, ...]]] =
 from numpy import cross, diagonal, matmul, outer, tensordot, trace
 
 __all__ = linalg_all.copy()
-__all__ += ['cross', 'diagonal', 'matmul', 'matrix_norm', 'outer',
-            'tensordot', 'trace', 'vector_norm']
+__all__ += ['cross', 'diagonal', 'matmul', 'matrix_norm', 'matrix_transpose',
+            'outer', 'svdvals', 'tensordot', 'trace', 'vecdot', 'vector_norm']