Repository: google/jax
Issue: linalg.svd and linalg.pinv fail for zero-size matrices

Hi all,

This issue follows from #5313 with PR #5316

`jnp.linalg.svd(jnp.zeros((0,0)))` and `jnp.linalg.pinv(jnp.zeros((0,0)))` both fail with runtime errors.

I will try and take a look if I find time, but I'm not sure how soon that will be I'm afraid :(



Changed files:
jax/_src/lax/linalg.py, tests/linalg_test.py, jax/_src/numpy/linalg.py

diff --git a/jax/_src/lax/linalg.py b/jax/_src/lax/linalg.py
index 68869af6c..2eefb47ee 100644
--- a/jax/_src/lax/linalg.py
+++ b/jax/_src/lax/linalg.py
@@ -1128,6 +1128,10 @@ def svd_impl(operand, full_matrices, compute_uv):
 def svd_translation_rule(c, operand, full_matrices, compute_uv):
   shape = c.get_shape(operand).dimensions()
   m, n = shape[-2:]
+  if m == 0 or n == 0:
+    return xla.lower_fun(_empty_svd, multiple_results=True)(
+      c, operand, full_matrices=full_matrices, compute_uv=compute_uv)
+
   u, s, v = xops.SVD(operand)
   permutation = list(range(len(shape)))
   permutation[-1], permutation[-2] = permutation[-2], permutation[-1]
@@ -1200,10 +1204,31 @@ def svd_jvp_rule(primals, tangents, full_matrices, compute_uv):
 
   return (s, U, Vt), (ds, dU, _H(dV))
 
+def _empty_svd(a, *, full_matrices, compute_uv):
+  batch_shape = a.shape[:-2]
+  m, n = a.shape[-2:]
+  s = jnp.empty(batch_shape + (0,), dtype=lax_internal._complex_basetype(a.dtype))
+  if not compute_uv:
+    return (s,)
+  if full_matrices:
+    size = max(m, n)
+    u = jnp.broadcast_to(jnp.eye(size, dtype=a.dtype), batch_shape + (size, size))
+  else:
+    u = jnp.empty(batch_shape + (m, n), dtype=a.dtype)
+  v = jnp.empty(batch_shape + (0, 0), dtype=a.dtype)
+  if m < n:
+    u, v = v, u
+  return s, u, v
+
 def _svd_cpu_gpu_translation_rule(gesvd_impl, c, operand, full_matrices, compute_uv):
+  shape = c.get_shape(operand).dimensions()
+  m, n = shape[-2:]
+  batch_dims = shape[:-2]
+
+  if m == 0 or n == 0:
+    return xla.lower_fun(_empty_svd, multiple_results=True)(
+      c, operand, full_matrices=full_matrices, compute_uv=compute_uv)
 
-  shape = c.get_shape(operand)
-  batch_dims = shape.dimensions()[:-2]
   s, u, vt, info = gesvd_impl(c, operand,
                               full_matrices=full_matrices,
                               compute_uv=compute_uv)
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
index c72774d7a..e4ec6be3c 100644
--- a/jax/_src/numpy/linalg.py
+++ b/jax/_src/numpy/linalg.py
@@ -303,7 +303,7 @@ def pinv(a, rcond=None):
   u, s, v = svd(a, full_matrices=False)
   # Singular values less than or equal to ``rcond * largest_singular_value``
   # are set to zero.
-  cutoff = rcond[..., jnp.newaxis] * jnp.amax(s, axis=-1, keepdims=True)
+  cutoff = rcond[..., jnp.newaxis] * jnp.amax(s, axis=-1, keepdims=True, initial=-jnp.inf)
   s = jnp.where(s > cutoff, s, jnp.inf)
   res = jnp.matmul(_T(v), jnp.divide(_T(u), s[..., jnp.newaxis]))
   return lax.convert_element_type(res, a.dtype)
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 81f8e3607..692c1a94b 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -513,8 +513,8 @@ class NumpyLinalgTest(jtu.JaxTestCase):
        "b": b, "m": m, "n": n, "dtype": dtype, "full_matrices": full_matrices,
        "compute_uv": compute_uv}
       for b in [(), (3,), (2, 3)]
-      for m in [2, 7, 29, 53]
-      for n in [2, 7, 29, 53]
+      for m in [0, 2, 7, 29, 53]
+      for n in [0, 2, 7, 29, 53]
       for dtype in float_types + complex_types
       for full_matrices in [False, True]
       for compute_uv in [False, True]))
@@ -529,7 +529,7 @@ class NumpyLinalgTest(jtu.JaxTestCase):
     # Norm, adjusted for dimension and type.
     def norm(x):
       norm = np.linalg.norm(x, axis=(-2, -1))
-      return norm / (max(m, n) * jnp.finfo(dtype).eps)
+      return norm / (max(1, m, n) * jnp.finfo(dtype).eps)
 
     a, = args_maker()
     out = jnp.linalg.svd(a, full_matrices=full_matrices, compute_uv=compute_uv)
@@ -773,7 +773,8 @@ class NumpyLinalgTest(jtu.JaxTestCase):
       {"testcase_name":
        "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
        "shape": shape, "dtype": dtype}
-      for shape in [(1, 1), (4, 4), (2, 70, 7), (2000, 7), (7, 1000), (70, 7, 2)]
+      for shape in [(1, 1), (4, 4), (2, 70, 7), (2000, 7), (7, 1000), (70, 7, 2),
+                    (2, 0, 0), (3, 0, 2), (1, 0)]
       for dtype in float_types + complex_types))
   def testPinv(self, shape, dtype):
     if (jnp.issubdtype(dtype, np.complexfloating) and