Basic DFT and CC examples fail on GPU

[frostedoyster]

Hello and thanks for this very useful library!

I've tried to run the basic DFT and CC examples, but I keep running into this error:

File ~/code/virtualenvs/jax/lib/python3.12/site-packages/pyscfad/backend/_jax/lax/linalg.py:165, in _eigh_gen_lowering(*args, **kwargs)
    [164](~/code/virtualenvs/jax/lib/python3.12/site-packages/pyscfad/backend/_jax/lax/linalg.py:164) def _eigh_gen_lowering(*args, **kwargs):
--> [165](~/code/virtualenvs/jax/lib/python3.12/site-packages/pyscfad/backend/_jax/lax/linalg.py:165)     raise NotImplementedError("Generalized eigh is only implemented for CPU.")

NotImplementedError: Generalized eigh is only implemented for CPU.

Does this mean that these examples can only be run on CPU? Thanks in advance

[fishjojo]

I haven't implemented the GPU backend for eigh.
One workaround is to replace the _eigh function in the hf.SCF class by the one in https://github.com/fishjojo/pyscfad/blob/main/pyscfad/ml/scf/hf.py.
But this may cause divergent gradients when there are degenerate orbitals.

I'll find a time to implement the GPU backend.