Repository: google/jax

Issue:
Implement JVP for SVD when full_matrices=True



Discussion:
Eqns 16-18 of [this note](https://j-towns.github.io/papers/svd-derivative.pdf) give formulae for the jvp. I think it’s probably also been implemented in Julia somewhere.
It’s been a while since I looked at this but just to warn anyone implementing this I think those formulae might be slightly more general than necessary because they hold for a ‘low rank svd’.

This means that one or both of the terms at the end of 16 and at the end of 18 will be identically zero, and therefore not worth wasting computation on.

You can infer which of the terms is zero by looking at the shapes of the inputs.

Resolving patch:
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -468,6 +468,32 @@ def svd_abstract_eval(operand, full_matrices, compute_uv):
     vt = operand
   return core.AbstractTuple((s, u, vt))
 
+def svd_jvp_rule(primals, tangents, full_matrices, compute_uv):
+  if full_matrices:
+    #TODO: implement full matrices case, documented here: https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf
+    raise NotImplementedError("Singular value decomposition JVP not implemented for full matrices")
+  A, = primals
+  dA, = tangents
+  s, U, Vt = svd_p.bind(A, full_matrices=False, compute_uv=True)
+
+  k = s.shape[-1]
+  Ut, V = np.conj(U).T, np.conj(Vt).T
+  s_dim = s[..., None, :]
+  dS = np.dot(np.dot(Ut, dA), V)
+  ds = np.diag(dS)
+  F = 1 / (np.square(s_dim) - np.square(s_dim.T) + np.eye(k)) - np.eye(k)
+  dSS = s_dim * dS
+  SdS = s_dim.T * dS
+  dU = np.dot(U, F * (dSS + dSS.T))
+  dV = np.dot(V, F * (SdS + SdS.T))
+
+  m, n = A.shape[-2], A.shape[-1]
+  if m > n:
+    dU = dU + np.dot(np.eye(m) - np.dot(U, Ut), np.dot(dA, V)) / s_dim
+  if n > m:
+    dV = dV + np.dot(np.eye(n) - np.dot(V, Vt), np.dot(np.conj(dA).T, U)) / s_dim
+  return core.pack((s, U, Vt)), core.pack((ds, dU, dV.T))
+
 def svd_cpu_translation_rule(c, operand, full_matrices, compute_uv):
   shape = c.GetShape(operand)
   dtype = shape.element_type().type
@@ -491,4 +517,5 @@ def svd_batching_rule(batched_args, batch_dims, full_matrices, compute_uv):
 svd_p.def_abstract_eval(svd_abstract_eval)
 xla.translations[svd_p] = svd_translation_rule
 xla.backend_specific_translations['cpu'][svd_p] = svd_cpu_translation_rule
+ad.primitive_jvps[svd_p] = svd_jvp_rule
 batching.primitive_batchers[svd_p] = svd_batching_rule


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -197,7 +197,7 @@ def testEighGradVectorComplex(self, shape, dtype, rng, lower, eps):
     # Assert rtol eigenvalue delta between perturbed eigenvectors vs new true eigenvalues.
     RTOL=1e-2
     assert onp.max(
-      onp.abs((onp.diag(onp.dot(onp.conj((v+dv).T), onp.dot(new_a,(v+dv)))) - new_w) / new_w)) < RTOL 
+      onp.abs((onp.diag(onp.dot(onp.conj((v+dv).T), onp.dot(new_a,(v+dv)))) - new_w) / new_w)) < RTOL
     # Redundant to above, but also assert rtol for eigenvector property with new true eigenvalues.
     assert onp.max(
       onp.linalg.norm(onp.abs(new_w*(v+dv) - onp.dot(new_a, (v+dv))), axis=0) /
@@ -284,6 +284,9 @@ def norm(x):
 
     self._CompileAndCheck(partial(np.linalg.svd, full_matrices=full_matrices, compute_uv=compute_uv),
                           args_maker, check_dtypes=True)
+    if not full_matrices:
+      svd = partial(np.linalg.svd, full_matrices=False)
+      jtu.check_jvp(svd, partial(jvp, svd), (a,), atol=1e-1 if FLAGS.jax_enable_x64 else jtu.ATOL)
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_fullmatrices={}".format(