suggestion on `jnp.square(jnp.linalg.norm(Rmnc, axis=3))`
https://code.itp.ac.cn/wanglei/hydrogen/-/blob/gaussian/src/mo.py#L64
-
give it a name, avoid recomputation of
jnp.square(jnp.linalg.norm(Rmnc, axis=3)) -
avoid using
square(norm), which is unfriendly to automatic differentiation
https://code.itp.ac.cn/wanglei/hydrogen/-/issues/12
Solved here https://code.itp.ac.cn/wanglei/hydrogen/-/issues/12#note_1598