def testSVD(self, b, m, n, dtype, full_matrices, compute_uv, hermitian):
    if (jnp.issubdtype(dtype, np.complexfloating) and
        jtu.device_under_test() == "tpu"):
      raise unittest.SkipTest("No complex SVD implementation")
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(b + (m, n), dtype)]

    # Norm, adjusted for dimension and type.
    def norm(x):
      norm = np.linalg.norm(x, axis=(-2, -1))
      return norm / (max(1, m, n) * jnp.finfo(dtype).eps)

    a, = args_maker()
    if hermitian:
      a = a + np.conj(T(a))
    out = jnp.linalg.svd(a, full_matrices=full_matrices, compute_uv=compute_uv,
                         hermitian=hermitian)
    if compute_uv:
      # Check the reconstructed matrices
      if full_matrices:
        k = min(m, n)
        if m < n:
          self.assertTrue(np.all(
              norm(a - np.matmul(out[1][..., None, :] * out[0], out[2][..., :k, :])) < 50))
        else:
          self.assertTrue(np.all(
              norm(a - np.matmul(out[1][..., None, :] * out[0][..., :, :k], out[2])) < 350))
      else:
        self.assertTrue(np.all(
          norm(a - np.matmul(out[1][..., None, :] * out[0], out[2])) < 350))

      # Check the unitary properties of the singular vector matrices.
      self.assertTrue(np.all(norm(np.eye(out[0].shape[-1]) - np.matmul(np.conj(T(out[0])), out[0])) < 15))
      if m >= n:
        self.assertTrue(np.all(norm(np.eye(out[2].shape[-1]) - np.matmul(np.conj(T(out[2])), out[2])) < 10))
      else:
        self.assertTrue(np.all(norm(np.eye(out[2].shape[-2]) - np.matmul(out[2], np.conj(T(out[2])))) < 20))

    else:
      self.assertTrue(np.allclose(np.linalg.svd(a, compute_uv=False), np.asarray(out), atol=1e-4, rtol=1e-4))

    self._CompileAndCheck(partial(jnp.linalg.svd, full_matrices=full_matrices, compute_uv=compute_uv),
                          args_maker)
    if not compute_uv:
      svd = partial(jnp.linalg.svd, full_matrices=full_matrices,
                    compute_uv=compute_uv)
      # TODO(phawkins): these tolerances seem very loose.
      if dtype == np.complex128:
        jtu.check_jvp(svd, partial(jvp, svd), (a,), rtol=1e-4, atol=1e-4, eps=1e-8)
      else:
        jtu.check_jvp(svd, partial(jvp, svd), (a,), rtol=5e-2, atol=2e-1)

    if jtu.device_under_test() == "tpu":
      raise unittest.SkipTest("TPU matmul does not have enough precision")
    # TODO(frederikwilde): Find the appropriate precision to use for this test on TPUs.

    if compute_uv and (not full_matrices):
      b, = args_maker()
      def f(x):
        u, s, v = jnp.linalg.svd(
          a + x * b,
          full_matrices=full_matrices,
          compute_uv=compute_uv)
        vdiag = jnp.vectorize(jnp.diag, signature='(k)->(k,k)')
        return jnp.matmul(jnp.matmul(u, vdiag(s)), v).real
      _, t_out = jvp(f, (1.,), (1.,))
      if dtype == np.complex128:
        atol = 1e-13
      else:
        atol = 5e-4
      self.assertArraysAllClose(t_out, b.real, atol=atol)