Flax 基础知识

• JAX 101
• Flax 基础
• 管理参数和状态
• setup vs compact
• 处理 Flax 模块参数
• Flax 中的随机性和 PRNG