TPU-Friendly Kolmogorov-Arnold Networks for Neural QMC: A 330x Compression of FermiNet Using JAX

[TL;DR] What if you could run Kolmogorov-Arnold Networks (KANs) efficiently on TPUs without suffering from memory fragmentation and conditional branching? Note for ML Engineers: While this project tac