Document if differentiability of implementation is dependent on hardware

[unalmis]

On my machine, eigh_tridagonal is reverse differentiable on cpu. On another machine, it is not reverse differentiable on gpu, causing unexpected nan leaks. It would be nice if the documentation of a function mentions when the differentiability of the current implementation is dependent on hardware.

[dfm]

Can you provide a minimal, reproducible example that shows the problem, and defines clearly what you mean by "not reverse differentiable". In other words, do you see an error, or does it seem like there are different numerics on different platforms which lead to this inconsistency? Thanks.

[unalmis]

A function that depends on the output of eigh_tridiagonal's eigenvalues gives a nan when jax.grad used on gpu. I believe this is expected due to current implementation

[dfm]

Like I suggested above, can you provide me with a function that I can run that demonstrates this and a pointer to where that belief is coming from? I'm not very familiar with this specific function, but I'm happy to try and debug or suggest where the docs should be updated if you can give me some more info!

[hawkinsp]

The structure of the spectrum is probably important here. In general eigendecomposition is not differentiable if, e.g., you have repeated eigenvalues. And it's possible that it's very sensitive to small numerical differences, e.g., between hardware.

[PhilipVinc]

Just a minor comment: I think eigendecomposition is in general differentiable (at least as long as the matrix is non singular).

For the case of degenerate spectrum, there are some complications but it is still 'defined'.
The problem is that the formula is more complicated and the operation is more costly, so no AD framework defaults to it.
For eigh I use this gist. For tridiagonal, you'd have to work your way through

https://gist.github.com/PhilipVinc/f92499294efb3787e123f89ace3c5b29