Repository: google/jax

Issue:
np.linalg.svd on TPU
JAX currently only supports SVD on the CPU or GPU and one gets the following error when attempting to run it on a TPU:

```
NotImplementedError: Singular value decomposition is only implemented on the CPU and GPU backends
```

Afaik the SVD is supported in tensorflow so there should exist XLA primitives for it. What needs to be done to add support in JAX?



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
@@ -854,8 +854,16 @@ def svd_impl(operand, full_matrices, compute_uv):
   return s, u, vt
 
 def svd_translation_rule(c, operand, full_matrices, compute_uv):
-  raise NotImplementedError(
-    "Singular value decomposition is only implemented on the CPU and GPU backends")
+  shape = c.get_shape(operand).dimensions()
+  m, n = shape[-2:]
+  u, s, v = xops.SVD(operand)
+  permutation = list(range(len(shape)))
+  permutation[-1], permutation[-2] = permutation[-2], permutation[-1]
+  vt = xops.Transpose(v, permutation)
+  if not full_matrices and m != n:
+    u = xops.SliceInDim(u, 0, min(m, n), stride=1, dimno=len(shape) - 1)
+    vt = xops.SliceInDim(vt, 0, min(m, n), stride=1, dimno=len(shape) - 2)
+  return xops.Tuple(c, [s, u, vt])
 
 def svd_abstract_eval(operand, full_matrices, compute_uv):
   if isinstance(operand, ShapedArray):


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -492,8 +492,10 @@ def testNorm(self, shape, dtype, ord, axis, keepdims, rng_factory):
       for full_matrices in [False, True]
       for compute_uv in [False, True]
       for rng_factory in [jtu.rand_default]))
-  @jtu.skip_on_devices("tpu")
   def testSVD(self, b, m, n, dtype, full_matrices, compute_uv, rng_factory):
+    if (jnp.issubdtype(dtype, np.complexfloating) and
+        jtu.device_under_test() == "tpu"):
+      raise unittest.SkipTest("No complex SVD implementation")
     rng = rng_factory(self.rng())
     _skip_if_unsupported_type(dtype)
     args_maker = lambda: [rng(b + (m, n), dtype)]
@@ -618,10 +620,12 @@ def testQrBatching(self, shape, dtype, rng_factory):
       for shape in [(1, 1), (4, 4), (2, 3, 5), (5, 5, 5), (20, 20), (5, 10)]
       for pnorm in [jnp.inf, -jnp.inf, 1, -1, 2, -2, 'fro']
       for dtype in float_types + complex_types))
-  @jtu.skip_on_devices("tpu")  # SVD is not implemented on the TPU backend
   @jtu.skip_on_devices("gpu")  # TODO(#2203): numerical errors
   def testCond(self, shape, pnorm, dtype):
     _skip_if_unsupported_type(dtype)
+    if (jnp.issubdtype(dtype, np.complexfloating) and
+        jtu.device_under_test() == "tpu"):
+      raise unittest.SkipTest("No complex SVD implementation")
 
     def gen_mat():
       # arr_gen = jtu.rand_some_nan(self.rng())
@@ -733,8 +737,10 @@ def args_maker():
       for shape in [(1, 1), (4, 4), (2, 70, 7), (2000, 7), (7, 1000), (70, 7, 2)]
       for dtype in float_types + complex_types
       for rng_factory in [jtu.rand_default]))
-  @jtu.skip_on_devices("tpu")  # SVD is not implemented on the TPU backend
   def testPinv(self, shape, dtype, rng_factory):
+    if (jnp.issubdtype(dtype, np.complexfloating) and
+        jtu.device_under_test() == "tpu"):
+      raise unittest.SkipTest("No complex SVD implementation")
     rng = rng_factory(self.rng())
     _skip_if_unsupported_type(dtype)
     args_maker = lambda: [rng(shape, dtype)]
@@ -742,10 +748,10 @@ def testPinv(self, shape, dtype, rng_factory):
     self._CheckAgainstNumpy(np.linalg.pinv, jnp.linalg.pinv, args_maker,
                             tol=1e-2)
     self._CompileAndCheck(jnp.linalg.pinv, args_maker)
-    # TODO(phawkins): 1e-1 seems like a very loose tolerance.
-    jtu.check_grads(jnp.linalg.pinv, args_maker(), 2, rtol=1e-1, atol=2e-1)
+    if jtu.device_under_test() != "tpu":
+      # TODO(phawkins): 1e-1 seems like a very loose tolerance.
+      jtu.check_grads(jnp.linalg.pinv, args_maker(), 2, rtol=1e-1, atol=2e-1)
 
-  @jtu.skip_on_devices("tpu")  # SVD is not implemented on the TPU backend
   def testPinvGradIssue2792(self):
     def f(p):
       a = jnp.array([[0., 0.],[-p, 1.]], jnp.float32) * 1 / (1 + p**2)
@@ -787,8 +793,10 @@ def testMatrixPower(self, shape, dtype, n, rng_factory):
       for shape in [(3, ), (1, 2), (8, 5), (4, 4), (5, 5), (50, 50)]
       for dtype in float_types + complex_types
       for rng_factory in [jtu.rand_default]))
-  @jtu.skip_on_devices("tpu")
   def testMatrixRank(self, shape, dtype, rng_factory):
+    if (jnp.issubdtype(dtype, np.complexfloating) and
+        jtu.device_under_test() == "tpu"):
+      raise unittest.SkipTest("No complex SVD implementation")
     rng = rng_factory(self.rng())
     _skip_if_unsupported_type(dtype)
     args_maker = lambda: [rng(shape, dtype)]