Natural Excited States to nan

[KingDamLiu]

Even in a small system like LiH with states=4, it is easy to nan in thousands to tens of thousands of steps. I am not sure why, can you help with this?

[dashu233]

While I'm not the owner of this repo, I met the same question as yours. In my experience, using float64 is helpful and solves almost all the nan issues. See https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision for details about using float64 in JAX.

[jsspencer]

Float64 is important in some cases but we did not need it for our excited state calculations. Other thoughts:

• Good pretraining can help
  • Carefully setting optimization hyperparameters can help (details of those we used are in our papers)
• Check if you are using TensorFloat32.

We just updated the code and the latest version might be worth trying. It is hard to know what the issue is without the exact setup and parameters you are using.

Finally, for research questions, please try sending email rather than (anonymous) github issues.