Add option to calculate the kinetic energy using scan instead of fori_loop.
Add option to calculate the kinetic energy using scan instead of fori_loop.

PiperOrigin-RevId: 407303101
Change-Id: I29febb6c07f47a92650e7186db7249757fb02b9f

deepmind/ferminet