MAINT: address review comments; convert linalg exceptions

Author: ev-br

torch_np/linalg.py

@@ -1,3 +1,5 @@
+import functools
+import math
 from typing import Sequence
 
 import torch
@@ -26,23 +28,29 @@ def _atleast_float_2(a, b):
     return a, b
 
 
-def _prod(iterable):
-    p = 1.0
-    for x in iterable:
-        p *= x
-    return p
+def linalg_errors(func):
+    @functools.wraps(func)
+    def wrapped(*args, **kwds):
+        try:
+            return func(*args, **kwds)
+        except torch._C._LinAlgError as e:
+            raise LinAlgError(*e.args)
+
+    return wrapped
 
 
 # ### Matrix and vector products ###
 
 
 @normalizer
+@linalg_errors
 def matrix_power(a: ArrayLike, n):
     a = _atleat_float_1(a)
     return torch.linalg.matrix_power(a, n)
 
 
 @normalizer
+@linalg_errors
 def multi_dot(inputs: Sequence[ArrayLike], *, out=None):
     return torch.linalg.multi_dot(inputs)
 
@@ -51,41 +59,46 @@ def multi_dot(inputs: Sequence[ArrayLike], *, out=None):
 
 
 @normalizer
+@linalg_errors
 def solve(a: ArrayLike, b: ArrayLike):
     a, b = _atleast_float_2(a, b)
     return torch.linalg.solve(a, b)
 
 
 @normalizer
+@linalg_errors
 def lstsq(a: ArrayLike, b: ArrayLike, rcond=None):
     a, b = _atleast_float_2(a, b)
     # NumPy is using gelsd: https://github.com/numpy/numpy/blob/v1.24.0/numpy/linalg/umath_linalg.cpp#L3991
-    return torch.linalg.lstsq(a, b, rcond=rcond, driver="gelsd")
+    # on CUDA, only `gels` is available though, so use it instead
+    driver = "gels" if a.is_cuda or b.is_cuda else "gelsd"
+    return torch.linalg.lstsq(a, b, rcond=rcond, driver=driver)
 
 
 @normalizer
+@linalg_errors
 def inv(a: ArrayLike):
     a = _atleast_float_1(a)
-    try:
-        result = torch.linalg.inv(a)
-    except torch._C._LinAlgError as e:
-        raise LinAlgError(*e.args)
+    result = torch.linalg.inv(a)
     return result
 
 
 @normalizer
+@linalg_errors
 def pinv(a: ArrayLike, rcond=1e-15, hermitian=False):
     a = _atleast_float_1(a)
     return torch.linalg.pinv(a, rtol=rcond, hermitian=hermitian)
 
 
 @normalizer
+@linalg_errors
 def tensorsolve(a: ArrayLike, b: ArrayLike, axes=None):
     a, b = _atleast_float_2(a, b)
     return torch.linalg.tensorsolve(a, b, dims=axes)
 
 
 @normalizer
+@linalg_errors
 def tensorinv(a: ArrayLike, ind=2):
     a = _atleast_float_1(a)
     return torch.linalg.tensorinv(a, ind=ind)
@@ -95,50 +108,44 @@ def tensorinv(a: ArrayLike, ind=2):
 
 
 @normalizer
+@linalg_errors
 def det(a: ArrayLike):
     a = _atleast_float_1(a)
     return torch.linalg.det(a)
 
 
 @normalizer
+@linalg_errors
 def slogdet(a: ArrayLike):
     a = _atleast_float_1(a)
     return torch.linalg.slogdet(a)
 
 
 @normalizer
+@linalg_errors
 def cond(x: ArrayLike, p=None):
     x = _atleast_float_1(x)
 
     # check if empty
     # cf: https://github.com/numpy/numpy/blob/v1.24.0/numpy/linalg/linalg.py#L1744
-    if x.numel() == 0 and _prod(x.shape[-2:]) == 0:
+    if x.numel() == 0 and math.prod(x.shape[-2:]) == 0:
         raise LinAlgError("cond is not defined on empty arrays")
 
     result = torch.linalg.cond(x, p=p)
 
-    # Convert nans to infs unless the original array had nan entries
+    # Convert nans to infs (numpy does it in a data-dependent way, depending on
+    # whether the input array has nans or not)
     # XXX: NumPy does this: https://github.com/numpy/numpy/blob/v1.24.0/numpy/linalg/linalg.py#L1744
-    # Do we want it? .any() synchronizes
-    """
-    nan_mask = torch.isnan(result)
-    if nan_mask.any():
-        nan_mask &= ~torch.isnan(x).any(axis=(-2, -1))  # XXX: any() w/ tuple axes
-        if result.ndim > 0:
-            result[nan_mask] = torch.inf
-        elif nan_mask:
-            result[()] = torch.inf
-    """
-
-    return result
+    return torch.where(torch.isnan(result), float("inf"), result)
 
 
 @normalizer
+@linalg_errors
 def matrix_rank(a: ArrayLike, tol=None, hermitian=False):
     a = _atleast_float_1(a)
 
     if a.ndim < 2:
-        return int(not (a == 0).all())
+        return int((a != 0).any())
 
     if tol is None:
         # follow https://github.com/numpy/numpy/blob/v1.24.0/numpy/linalg/linalg.py#L1885
@@ -150,6 +157,7 @@ def matrix_rank(a: ArrayLike, tol=None, hermitian=False):
 
 
 @normalizer
+@linalg_errors
 def norm(x: ArrayLike, ord=None, axis=None, keepdims=False):
     x = _atleast_float_1(x)
     result = torch.linalg.norm(x, ord=ord, dim=axis)
@@ -162,12 +170,14 @@ def norm(x: ArrayLike, ord=None, axis=None, keepdims=False):
 
 
 @normalizer
+@linalg_errors
 def cholesky(a: ArrayLike):
     a = _atleast_float_1(a)
     return torch.linalg.cholesky(a)
 
 
 @normalizer
+@linalg_errors
 def qr(a: ArrayLike, mode="reduced"):
     a = _atleast_float_1(a)
     result = torch.linalg.qr(a, mode=mode)
@@ -178,19 +188,22 @@ def qr(a: ArrayLike, mode="reduced"):
 
 
 @normalizer
+@linalg_errors
 def svd(a: ArrayLike, full_matrices=True, compute_uv=True, hermitian=False):
     a = _atleast_float_1(a)
+    if not compute_uv:
+        return torch.linalg.svdvals(a)
+
     # NB: ignore the hermitian= argument (no pytorch equivalent)
     result = torch.linalg.svd(a, full_matrices=full_matrices)
-    if not compute_uv:
-        result = result.S
     return result
 
 
 # ### Eigenvalues and eigenvectors ###
 
 
 @normalizer
+@linalg_errors
 def eig(a: ArrayLike):
     a = _atleast_float_1(a)
     w, vt = torch.linalg.eig(a)
@@ -203,12 +216,14 @@ def eig(a: ArrayLike):
 
 
 @normalizer
+@linalg_errors
 def eigh(a: ArrayLike, UPLO="L"):
     a = _atleast_float_1(a)
     return torch.linalg.eigh(a, UPLO=UPLO)
 
 
 @normalizer
+@linalg_errors
 def eigvals(a: ArrayLike):
     a = _atleast_float_1(a)
     result = torch.linalg.eigvals(a)
@@ -219,6 +234,7 @@ def eigvals(a: ArrayLike):
 
 
 @normalizer
+@linalg_errors
 def eigvalsh(a: ArrayLike, UPLO="L"):
     a = _atleast_float_1(a)
     return torch.linalg.eigvalsh(a, UPLO=UPLO)

torch_np/tests/numpy_tests/linalg/test_linalg.py

@@ -705,7 +705,6 @@ def test_basic_nonsvd(self):
         assert_almost_equal(linalg.cond(A, -1), 0.5)
         assert_almost_equal(linalg.cond(A, 'fro'), np.sqrt(265 / 12))
 
-    @pytest.mark.xfail(reason="nan -> inf, as numpy does")
     def test_singular(self):
         # Singular matrices have infinite condition number for
         # positive norms, and negative norms shouldn't raise
@@ -747,7 +746,6 @@ def test_nan(self):
                 assert_(not np.isnan(c[0]))
                 assert_(not np.isnan(c[2]))
 
-    @pytest.mark.xfail(reason="nan -> inf, as numpy does")
     def test_stacked_singular(self):
         # Check behavior when only some of the stacked matrices are
         # singular