Repository: google/jax

Issue:
[feature request] add support for numpy.linalg.svd and scipy.linalg.svd
Hi,

Considering there has been recent work going on w.r.t. improving the linalg modules with the addition of Cholesky and QR decompositions and solve methods, I would like to see SVD added to this list.

I don’t mind working on this with some guidance from the developers, if they think this is worth adding. Thank you.


Discussion:

Resolving patch:
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -42,6 +42,13 @@ def qr(x, full_matrices=True):
   q, r = qr_p.bind(x, full_matrices=full_matrices)
   return q, r
 
+def svd(x, full_matrices=True, compute_uv=True):
+  s, u, v = svd_p.bind(x, full_matrices=full_matrices, compute_uv=compute_uv)
+  if compute_uv:
+    return u, s, v
+  else:
+    return s
+
 def triangular_solve(a, b, left_side=False, lower=False, transpose_a=False,
                      conjugate_a=False):
   return triangular_solve_p.bind(
@@ -316,3 +323,49 @@ def qr_jvp_rule(primals, tangents, full_matrices):
 qr_p.def_abstract_eval(qr_abstract_eval)
 xla.translations[qr_p] = qr_translation_rule
 ad.primitive_jvps[qr_p] = qr_jvp_rule
+
+
+# Singular value decomposition
+
+def svd_impl(operand, full_matrices, compute_uv):
+  s, u, vt = xla.apply_primitive(svd_p, operand, full_matrices=full_matrices, compute_uv=compute_uv)
+  return core.pack((s, u, vt))
+
+def svd_translation_rule(c, operand, full_matrices, compute_uv):
+  raise NotImplementedError(
+    "Singular value decomposition is only implemented on the CPU backend")
+
+def svd_abstract_eval(operand, full_matrices, compute_uv):
+  if isinstance(operand, ShapedArray):
+    if operand.ndim < 2:
+      raise ValueError("Argument to singular value decomposition must have ndims >= 2")
+
+    batch_dims = operand.shape[:-2]
+    m = operand.shape[-2]
+    n = operand.shape[-1]
+    s = ShapedArray(batch_dims + (min(m, n),), operand.dtype)
+    u = ShapedArray(batch_dims + (m, m if full_matrices else min(m, n)), operand.dtype)
+    vt = ShapedArray(batch_dims + (n if full_matrices else min(m, n), n), operand.dtype)
+  else:
+    s = operand
+    u = operand
+    vt = operand
+  return core.AbstractTuple((s, u, vt))
+
+def svd_cpu_translation_rule(c, operand, full_matrices, compute_uv):
+  shape = c.GetShape(operand)
+  dtype = shape.element_type().type
+  if len(shape.dimensions()) == 2 and dtype in {np.float32, np.float64}:
+    out = lapack.jax_gesdd(c, operand, full_matrices=full_matrices, compute_uv=compute_uv)
+    return c.Tuple(c.GetTupleElement(out, 0),
+                   c.GetTupleElement(out, 1),
+                   c.GetTupleElement(out, 2))
+  else:
+    raise NotImplementedError(
+        "Only unbatched singular value decomposition for real matrices is implemented on CPU")
+
+svd_p = Primitive('svd')
+svd_p.def_impl(svd_impl)
+svd_p.def_abstract_eval(svd_abstract_eval)
+xla.translations[svd_p] = svd_translation_rule
+xla.backend_specific_translations['Host'][svd_p] = svd_cpu_translation_rule
diff --git a/jax/numpy/linalg.py b/jax/numpy/linalg.py
--- a/jax/numpy/linalg.py
+++ b/jax/numpy/linalg.py
@@ -55,6 +55,13 @@ def cholesky(a):
   return lax_linalg.cholesky(a)
 
 
+@_wraps(onp.linalg.svd)
+def svd(a, full_matrices=True, compute_uv=True):
+  warnings.warn(_EXPERIMENTAL_WARNING)
+  a = _promote_arg_dtypes(np.asarray(a))
+  return lax_linalg.svd(a, full_matrices, compute_uv)
+
+
 @_wraps(onp.linalg.slogdet)
 def slogdet(a):
   a = _promote_arg_dtypes(np.asarray(a))
diff --git a/jax/scipy/linalg.py b/jax/scipy/linalg.py
--- a/jax/scipy/linalg.py
+++ b/jax/scipy/linalg.py
@@ -40,6 +40,14 @@ def cholesky(a, lower=False, overwrite_a=False, check_finite=True):
   return l if lower else np.conj(l.T)
 
 
+@_wraps(scipy.linalg.svd)
+def svd(a, full_matrices=True, compute_uv=True, overwrite_a=False, check_finite=True, lapack_driver='gesdd'):
+  warnings.warn(_EXPERIMENTAL_WARNING)
+  del overwrite_a, check_finite, lapack_driver
+  a = np_linalg._promote_arg_dtypes(np.asarray(a))
+  return lax_linalg.svd(a, full_matrices, compute_uv)
+
+
 @_wraps(scipy.linalg.det)
 def det(a, overwrite_a=False, check_finite=True):
   warnings.warn(_EXPERIMENTAL_WARNING)


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -134,6 +134,54 @@ def norm(x):
     self._CompileAndCheck(partial(np.linalg.eigh, UPLO=uplo), args_maker,
                           check_dtypes=True)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_n={}_full_matrices={}_compute_uv={}".format(
+          jtu.format_shape_dtype_string((m, n), dtype), full_matrices, compute_uv),
+       "m": m, "n": n, "dtype": dtype, "full_matrices": full_matrices,
+       "compute_uv": compute_uv, "rng": rng}
+      for m in [2, 7, 29, 53]
+      for n in [2, 7, 29, 53]
+      for dtype in float_types()
+      for full_matrices in [False, True]
+      for compute_uv in [False, True]
+      for rng in [jtu.rand_default()]))
+  @jtu.skip_on_devices("gpu", "tpu")
+  def testSVD(self, m, n, dtype, full_matrices, compute_uv, rng):
+    if not hasattr(lapack, "jax_gesdd"):
+      self.skipTest("No singular value decomposition implementation available")
+    args_maker = lambda: [rng((m, n), dtype)]
+
+    # Norm, adjusted for dimension and type.
+    def norm(x):
+      n = onp.linalg.norm(x, axis=(-2, -1))
+      return n / (max(m, n) * onp.finfo(dtype).eps)
+
+    a, = args_maker()
+    out = np.linalg.svd(a, full_matrices=full_matrices, compute_uv=compute_uv)
+
+    if compute_uv:
+      # Check the reconstructed matrices
+      if full_matrices:
+        k = min(m, n)
+        if m < n:
+          self.assertTrue(onp.all(norm(a - onp.matmul(out[1] * out[0], out[2][:k, :])) < 50))
+        else:
+          self.assertTrue(onp.all(norm(a - onp.matmul(out[1] * out[0][:, :k], out[2])) < 50))
+      else:
+          self.assertTrue(onp.all(norm(a - onp.matmul(out[1] * out[0], out[2])) < 50))
+
+      # Check the unitary properties of the singular vector matrices.
+      self.assertTrue(onp.all(norm(onp.eye(out[0].shape[1]) - onp.matmul(T(out[0]), out[0])) < 5))
+      if m >= n:
+        self.assertTrue(onp.all(norm(onp.eye(out[2].shape[1]) - onp.matmul(T(out[2]), out[2])) < 5))
+      else:
+        self.assertTrue(onp.all(norm(onp.eye(out[2].shape[0]) - onp.matmul(out[2], T(out[2]))) < 5))
+
+    else:
+      self.assertTrue(onp.allclose(onp.linalg.svd(a, compute_uv=False), onp.asarray(out)))
+
+    self._CompileAndCheck(partial(np.linalg.svd, full_matrices=full_matrices, compute_uv=compute_uv),
+                          args_maker, check_dtypes=True)
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_fullmatrices={}".format(