Possible bug in eigh derivative (for non-degenerate eigenvalues).

[sschoenholz]

We stumbled across an error in the gradient when using eigh without any degenerate eigenvalues. Here is a simple repro (tested in Google colab with JAX version 0.3.4).

import jax.numpy as jnp
import jax.test_util as jtu

def fn(x):
  A = x[:, None] * x[None, :]
  _, U = jnp.linalg.eigh(A)
  return U @ x

x = jnp.array([1.0, 2.0, 3.0])
print(x[:, None] * x[None, :])  # Symmetric matrix.

print(fn(x))  # Reasonable output 

jtu.check_grads(fn, (x,), 1)

"""
AssertionError: 
Not equal to tolerance rtol=0.006, atol=0.006
JVP tangent
Mismatched elements: 3 / 3 (100%)
Max absolute difference: 5823.4707
Max relative difference: 1.0003492
 x: array([ 1.707243,  0.97609 , -1.043212], dtype=float32)
 y: array([ 2682.228 , -5822.4946,  2987.9426], dtype=float32)
"""

[YouJiacheng]

Why you said that a rank-1 3×3 matrix is "without any degenerate eigenvalues"?
For the title "eigh derivative for non-degenerate eigenvalues", I guess you might want to have well defined gradient with the part of eigenvalues which are not degenerate?
It seems that only the case that all eigenvalues is not degenerate is implemented.
https://github.com/google/jax/blob/44006c7e5c7032a5560a9bd62394034ea42ce3da/jax/_src/lax/linalg.py#L550-L558

[jakevdp]

Hi - this looks like it does have degenerate eigenvalues (the eigenvalues are [0, 0, 14]) I think that in this case the gradient is not well defined - see the discussion in #669

[sschoenholz]

Oops you're right! Hmm, back to the repro drawing board... this case is definitely degenerate.