Fix some issues with the linalg wrapping

Author: asmeurer

array_api_compat/common/_linalg.py

@@ -59,7 +59,7 @@ def svd(x: ndarray, /, xp, *, full_matrices: bool = True) -> SVDResult:
 def cholesky(x: ndarray, /, xp, *, upper: bool = False) -> ndarray:
     L = xp.linalg.cholesky(x)
     if upper:
-        return matrix_transpose(L)
+        return get_xp(xp)(matrix_transpose)(L)
     return L
 
 # The rtol keyword argument of matrix_rank() and pinv() is new from NumPy.
@@ -158,7 +158,8 @@ def diagonal(x: ndarray, /, xp, *, offset: int = 0) -> ndarray:
 def trace(x: ndarray, /, xp, *, offset: int = 0) -> ndarray:
     return xp.asarray(xp.trace(x, offset=offset, axis1=-2, axis2=-1))
 
-__all__ = ['cross', 'diagonal', 'matmul', 'cholesky', 'matrix_rank', 'pinv',
-           'matrix_norm', 'matrix_transpose', 'outer', 'svdvals',
-           'tensordot', 'trace', 'vecdot', 'vector_norm', 'EighResult',
-           'QRResult', 'SlogdetResult', 'SVDResult']
+__all__ = ['cross', 'matmul', 'outer', 'tensordot', 'EighResult',
+           'QRResult', 'SlogdetResult', 'SVDResult', 'eigh', 'qr', 'slogdet',
+           'svd', 'cholesky', 'matrix_rank', 'pinv', 'matrix_norm',
+           'matrix_transpose', 'svdvals', 'vecdot', 'vector_norm', 'diagonal',
+           'trace']

array_api_compat/cupy/linalg.py

@@ -14,19 +14,27 @@
 import cupy as cp
 
 cross = get_xp(cp)(_linalg.cross)
-diagonal = get_xp(cp)(_linalg.diagonal)
 matmul = get_xp(cp)(_linalg.matmul)
+outer = get_xp(cp)(_linalg.outer)
+tensordot = get_xp(cp)(_linalg.tensordot)
+EighResult = _linalg.EighResult
+QRResult = _linalg.QRResult
+SlogdetResult = _linalg.SlogdetResult
+SVDResult = _linalg.SVDResult
+eigh = get_xp(cp)(_linalg.eigh)
+qr = get_xp(cp)(_linalg.qr)
+slogdet = get_xp(cp)(_linalg.slogdet)
+svd = get_xp(cp)(_linalg.svd)
 cholesky = get_xp(cp)(_linalg.cholesky)
 matrix_rank = get_xp(cp)(_linalg.matrix_rank)
 pinv = get_xp(cp)(_linalg.pinv)
 matrix_norm = get_xp(cp)(_linalg.matrix_norm)
 matrix_transpose = get_xp(cp)(_linalg.matrix_transpose)
-outer = get_xp(cp)(_linalg.outer)
 svdvals = get_xp(cp)(_linalg.svdvals)
-tensordot = get_xp(cp)(_linalg.tensordot)
-trace = get_xp(cp)(_linalg.trace)
 vecdot = get_xp(cp)(_linalg.vecdot)
 vector_norm = get_xp(cp)(_linalg.vector_norm)
+diagonal = get_xp(cp)(_linalg.diagonal)
+trace = get_xp(cp)(_linalg.trace)
 
 __all__ = linalg_all + _linalg.__all__
 

array_api_compat/numpy/linalg.py

@@ -7,19 +7,27 @@
 import numpy as np
 
 cross = get_xp(np)(_linalg.cross)
-diagonal = get_xp(np)(_linalg.diagonal)
 matmul = get_xp(np)(_linalg.matmul)
+outer = get_xp(np)(_linalg.outer)
+tensordot = get_xp(np)(_linalg.tensordot)
+EighResult = _linalg.EighResult
+QRResult = _linalg.QRResult
+SlogdetResult = _linalg.SlogdetResult
+SVDResult = _linalg.SVDResult
+eigh = get_xp(np)(_linalg.eigh)
+qr = get_xp(np)(_linalg.qr)
+slogdet = get_xp(np)(_linalg.slogdet)
+svd = get_xp(np)(_linalg.svd)
 cholesky = get_xp(np)(_linalg.cholesky)
 matrix_rank = get_xp(np)(_linalg.matrix_rank)
 pinv = get_xp(np)(_linalg.pinv)
 matrix_norm = get_xp(np)(_linalg.matrix_norm)
 matrix_transpose = get_xp(np)(_linalg.matrix_transpose)
-outer = get_xp(np)(_linalg.outer)
 svdvals = get_xp(np)(_linalg.svdvals)
-tensordot = get_xp(np)(_linalg.tensordot)
-trace = get_xp(np)(_linalg.trace)
 vecdot = get_xp(np)(_linalg.vecdot)
 vector_norm = get_xp(np)(_linalg.vector_norm)
+diagonal = get_xp(np)(_linalg.diagonal)
+trace = get_xp(np)(_linalg.trace)
 
 __all__ = linalg_all + _linalg.__all__