Incorrect SVD using `jax.linalg.svd` - Python jax

I have a code that uses singular value decomposition (SVD) to obtain an operator and then compute the Jacobian for it using jacfwd. But I was getting NaNs so I tried to test the implementation of SVD in Jax for complex matrices. The following snippet shows one instance where the Jax implementation of SVD fails with errors of order 1e-7. With the suggestion by @hawkinsp I tried using 64-bit precision and still see an error. This error somehow depends on the singular values and it is possible to get larger errors in some cases.

import numpy as np
from scipy.stats import unitary_group

from jax.config import config
from jax import numpy as jnp
import jax


print("Jax version", jax.__version__)
config.update("jax_enable_x64", True)


np.random.seed(42)

N = 15

# Construct an operator using two random Unitaries
U = unitary_group.rvs(N)
V = unitary_group.rvs(N)
S = np.zeros((N, N))
np.fill_diagonal(S, np.random.uniform(0.5, 1., size = (N, N)))

op = np.dot(U*S, jnp.conjugate(V).T).astype(jnp.complex128)

u, s, v = jnp.linalg.svd(op)
assert(np.allclose(op, jnp.dot(u*s, v)))

which results in

Jax version 0.1.70
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-79-babe60e251fc> in <module>
     25 
     26 u, s, v = jnp.linalg.svd(op)
---> 27 assert(np.allclose(op, jnp.dot(u*s, v)))

AssertionError: 

Numpy however gives the correct decomposition with

u, s, v = np.linalg.svd(op)
np.allclose(op, np.dot(u*s, v))

If I compare the error in the reconstructed matrix after SVD, it is of the order 1e-7 using Jax:

Numpy however gives a much better result with errors of the order 1e-16:

Is it possible to test the svd function better (or with lower tolerances) since right now the tests are with tol=1e-4 and rtol=1e-4 in https://github.com/google/jax/blob/master/tests/linalg_test.py#L584

@hawkinsp makes a comment that the tests are very loose. I am wondering if this is leading to issues in my code since I have a forward function that takes SVD and then removes negative singular values and computes a Jacobian for the resulting matrix. I get NaNs for such a computation.

Asked Oct 11 '21 14:10

quantshah