Grad calculation when “jax.scipy.linalg.eigh()” is involved

[miaoL1]

I've encountered a challenge calculating the gradient value after eigenvectors were computed.
I have simplified my problem in code as below:

import jax
import jax.numpy as jnp
def parameter(theta):
    H = jnp.ones([8,8]) + theta*jnp.eye(8)
    Q = jax.scipy.linalg.eigh(H)[1]
    return jnp.abs(jnp.linalg.slogdet(Q)[1])
gradient_fn = jax.grad(parameter)
theta_value = 10.0
gradient = gradient_fn(theta_value)
print("Gradient with respect to theta:", gradient)

No matter what "theta value " is given, the result is always "nan".
I think understanding of this problem might help me solve my difficulty.
I'd appreciate it if anyone can help.

[jakevdp]

Hi! This is a fundamental issue with autodiff of eigenvalue problems: the gradient in the presence of degenerate eigenvalues is mathematically ill-defined (see e.g. #669 and #4646), and so NaN is a reasonable output for an output with no well-defined value.

If you replace your function with one that uses a non-degenerate input matrix, you should see a non-NaN output. Try this, for example:

def parameter(theta):
    key = jax.random.key(0)
    H = jax.random.normal(key, (8, 8)) + theta*jnp.eye(8)
    Q = jax.scipy.linalg.eigh(H)[1]
    return jnp.abs(jnp.linalg.slogdet(Q)[1])

[hawkinsp]

Another comment: do you actually want the derivative of the eigendecomposition itself (not well-defined mathematically?) Or do you intend to do some larger computation with that eigendecomposition? Perhaps you can write a custom derivative rule for that larger thing?