Flax fundamentals# JAX 101 Flax Basics Managing Parameters and State setup vs compact Dealing with Flax Module arguments Randomness and PRNGs in Flax
