迁移# 从 NNX 0.10 迁移到 NNX 0.11 在 NNX 变换中使用 Rngs 加载带 RNG 的检查点 优化器更新 包含 NNX 对象的 Pytrees 从 Flax Linen 迁移到 Flax NNX 基本的 Module 定义 变量创建 训练步骤和编译 集合和变量类型 使用多种方法 变换 扫描层 在 Flax NNX 中使用 TrainState 从 Haiku 迁移到 Flax NNX 基本的模块定义 变量创建 训练步骤和编译 处理非参数状态 使用多种方法 变换 扫描层 顶层 Haiku 函数与顶层 Flax 模块
