指南# Flax 基础知识 JAX 101 Flax 基础 管理参数和状态 setup 与 compact 处理 Flax 模块参数 Flax 中的随机性和 PRNG 数据预处理 处理整个数据集 加载数据集 训练技巧 批量归一化 丢弃 学习率调度 迁移学习 保存和加载检查点 并行训练 在多个设备上进行集成 在多个设备上扩展 Flax 模块 模型检查 模型手术 提取中间值 提取中间值的梯度 转换和升级 从 Haiku 迁移到 Flax 将 PyTorch 模型转换为 Flax 将检查点迁移到 Orbax 将代码库升级到 Optax 将代码库升级到 Linen RNNCellBase 升级指南 迁移到常规字典 量化 使用 FP8 的用户指南 核心部分 🔪 flax.linen.Dropout 层和随机性