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

How leveraging JAX (jnp.einsum and lax.scan) transformed branching KAN architectures into TPU-friendly dense GEMMs, unlocking "Glass-Box" interpretability.