SVD Non-Uniqueness

[EricaCMitchell]

Hello,

I have been working with JAX and have run into the problem that I need the vectors that span the null space of a SVD. I have seen that on #508 the SVD derivatives for a matrix (m!=n) are undefined for the full-rank case due to non-uniqueness. Is there no workaround to compute this?

What I'm trying to do is get these null space vectors so that two conditions are fulfilled:

1. For the non-square overlap matrix (S12), where S12 = U @ Sigma @ Vt,
   U @ Sigma @ Vt @ Vt[nullspace, :] == 0
   and its derivatives equal 0 as well
2. For the square overlap matrix (S22),
   Vt[nullspace, :].T @ S22 @ Vt[nullspace, :] == 1
   ans its derivatives equal 0

The code below makes it so the basis functions of space 2 exactly span the basis functions of space 1.

 # Compute the overlap matrix between space 1 and space 2
S12 = compute_f12_oeints(geom, basis_set, complementary_aux_basis_set, xyz_path, deriv_order, options, True)

# Compute the eigenvectors and eigenvalues of C2.T @ S12.T @ S12 @ C2
S22 = jnp.dot(S12.T, S12)
CTC = C2.T @ S22 @ C2
Sigma_squared, V = jnp.linalg.eigh(CTC)

nullspace = jnp.where(Sigma_squared < 1.0e-6, True, False)

V_N = V[:, nullspace]

# Half-backtransform to space 2
C_nullspace = jnp.dot(C2, V_N)

The two conditions are fulfilled when the derivative is not taken, but the second condition fails to have the derivatives equal 0.

[EricaCMitchell]

I see that in #508 the autodiff formulation of the SVD for case m <= n and m > n are given in Giles2008. Although numerical instabilities may still occur due to the inversion of the singular values and degenerate singular values, it seems to me that an implementation of SVD full_matrices=True is possible with the same caveats as other eigendecomposition methods.

In a somewhat recent GitHub project, the SVD derivative is done using the work of Townsend and the works of others to compute the derivative of the SVD for the cases:

1. Real valued partial SVD
2. Real valued full SVD
3. Complex Input SVD

Although the nonuniqueness of the null space vectors of U or VT is present, I don't understand why the full_matrices=True is not able to be implemented in JAX natively.