Fix Unit test failure for JAX/Flax version update#262

susanbao wants to merge 9 commits intomainAI-Hypercomputer/maxdiffusion:mainfrom

sanbao/qwix_errorAI-Hypercomputer/maxdiffusion:sanbao/qwix_errorCopy head branch name to clipboard