jax.numpy.linalg.svd & numpy.linalg.svd => different conventions?

[jecampagne]

Hello,
I experience a rather subtle/nasty difference between jax.numpy.linalg.svd & numpy.linalg.svd.

(nb. after first edition, I run also jax.scipy.linalg.svd & jax.lax.linalg.svd they both agree to jax.numpy.linalg.svd, so they are different than numpy.linalg.svd)
(nb new add: numpy.linalg.svd agrees with Mathematica convention)

I'm running

jax                       0.3.5                    pypi_0    pypi
jaxlib                    0.3.5+cuda11.cudnn82          pypi_0    pypi

Here is a snippet that reproduces the problem:

import jax
import jax.numpy as jnp
import numpy as np

X= jax.random.normal(jax.random.PRNGKey(42),shape=(5,3))

print("X:",X,"\n\n")

cov_jax = jnp.cov(X.T)
a_jax,b_jax,c_jax = jnp.linalg.svd(cov_jax)

cov_np = np.cov(X.T)
a_np, b_np, c_np = np.linalg.svd(cov_np)

print("cov:",cov_jax,"\n\n",cov_np, "\n\n")

print("a:",a_jax,"\n\n",a_np,"\n\n")

print("b:",b_jax,"\n\n",b_np,"\n\n")

print("c:",c_jax,"\n\n",c_np)

which leads to

X: [[-0.716899   -0.20865498 -2.5713923 ]
 [ 1.0337092  -1.4035789   1.3873734 ]
 [-0.8396519   0.3010434   0.1421263 ]
 [-1.7631724  -1.6755074   0.31390068]
 [ 0.5912831   0.5325395  -0.9133108 ]] 


cov: [[ 1.292871    0.22079568  0.37710828]
 [ 0.22079568  0.9976356  -0.7964721 ]
 [ 0.37710828 -0.7964721   2.2377398 ]] 

 [[ 1.292871    0.22079572  0.37710825]
 [ 0.22079572  0.99763561 -0.7964721 ]
 [ 0.37710825 -0.7964721   2.23773968]] 


a: [[ 0.18085091  0.8960104  -0.40553415]
 [-0.40301841  0.44364002  0.8004744 ]
 [ 0.8971449   0.01867123  0.4413415 ]] 

 [[-0.18085096 -0.89601054 -0.40553427]
 [ 0.40301855 -0.44364005  0.80047458]
 [-0.8971449  -0.01867124  0.4413416 ]] 


b: [2.6715527 1.4100511 0.4466418] 

 [2.67155301 1.41005143 0.44664184] 


c: [[ 0.18085095 -0.40301847  0.89714485]
 [ 0.89601046  0.44364005  0.01867122]
 [-0.40553427  0.8004746   0.44134155]] 

 [[-0.18085096  0.40301855 -0.8971449 ]
 [-0.89601054 -0.44364005 -0.01867124]
 [-0.40553427  0.80047458  0.4413416 ]]

As you see, the "a" and "c" of the SVD have sign differences between Numpy and Jax.Numpy...

Is-it foreseen? Does it is well documented/advertised (I do not find mention of this in the doc) ?

Thanks.

[jakevdp]

Hi - in general, there is no preferred sign for eigenvectors of an svd. Even for numpy itself, different systems will return eigenvectors of different signs for the same input (see e.g. numpy/numpy#21313). As such, I would say what you're seeing is expected behavior.

[jakevdp]

Right, but again I think this is to be expected. Singular vectors are only defined up to a sign, and so different svd implementations may return different signs and still be correct.

Keep in mind also this is not really something JAX can control: for example, if you compile JAX against different LAPACK implementations (e.g. OpenBlas vs. MKL), you may get different sign conventions even on the same hardware.

Can you say more about why the signs of singular vectors are important in your application?

[jecampagne]

This is the piece of code that an third party uses that I was transfering to JAX, and find the above mentonned pb switching to GPU.
(the theta below is the X that I have used in the snippet)

        # calculate the covariance of the inputs
        cov = np.cov(theta.T)
        # calculate the Singular Value Decomposition
        a, b, c = np.linalg.svd(cov)
        m_diag = np.diag(1.0 / np.sqrt(b))
        # the transformation matrix
        mu_matrix = np.dot(m_diag, c)
        # calculate the transformed input parameters
        theta_trans = np.dot(mu_matrix,theta.T).T

[jakevdp]

Thanks, but why is it problematic that theta_trans can differ by a sign? You have the vectors a which correctly encode how they are related to the inputs.

Again, this issue also exists without consideration of JAX, because this numpy code may behave differently depending on the platform and which particular LAPACK library numpy is compiled against.

[jecampagne]

This is not just a global sign...

[jakevdp]

I see - I understand that you're combining singular vectors in a way that the resulting output is dependent on the sign convention chosen by your SVD implementation. But given that the SVD by definition is free to return either sign, why is it important? It seems like creating an algorthm whose results are dependent on an undefined convention indicates an issue with that algorithm, not an issue with the convention.

Again, this has nothing to do with JAX per se: if you run that numpy code on different backends, you will get different outputs. Perhaps you can more carefully define your algorithm to account for the fact that the sign convention for a given SVD implementations is not defined?

[YouJiacheng]

I think you should use the sign of the sequential principal minor determinant to choose the sign of eigenvectors, in order to have a portable implementation.

[jecampagne]

Hi, thanks for your interest. Can you elaborate a little your last hint?
The piece of code I have given above, is not my algo but, I tink your right that it should be independant of the SVD implementation. So thanks for your help.

[YouJiacheng]

Compute the sign of det(c[:i, :i]) for i = 1, 2, 3...
Then you can flip the sign of eigenvectors to satisfy the sign convention.
It seems that jax use +++... while numpy use -++...?
BTW, since cov is PSD, why not use eigh?