用途：为 JAX、Flax、Optax 训练任务产出一份可实现的训练配方，覆盖从实验定义到多设备扩展的完整流程。
主要能力：先收集任务、数据集、损失与指标，目标设备（CPU、单 GPU、多 GPU、TPU），模型框架（Flax/Linen、Equinox 或原生 JAX），随机性需求（dropout、采样、增强），状态需求（batch norm、优化器状态、EMA、检查点）。随后规划 PRNG 键的归属与拆分、模型与状态边界、初始化模型与优化器，编写纯 train_step，并在未 jit 版本可用后再逐步加入 jit、vmap、pmap，最后补充检查点与恢复测试。
使用方式：按工作流推进——定义实验与最小数据批次；规划 PRNG 策略；定义模型/状态边界；初始化；写纯 train_step；先不加 jit 验证；单设备正确后再加 vmap/pmap；加检查点与恢复测试；执行首次运行安全检查；从极小批次逐步扩展到完整训练。具体配置与清单细节见 references/training-recipe-details.md。
JAX 特定规则：绝不重用 PRNG 键；保持状态显式，不把可变状态藏在闭包中；避免 jit 内的 Python 副作用；把 jit 下的形状变化视为设计决策；先不用 jit 调试，再逐个重新启用变换；多设备训练时有意复制状态。
常见失败模式：PRNG 键重用导致 dropout/采样重复；Python 控制流或副作用引发 tracer 错误；batch norm/dropout 状态未与参数分离；vmap 或 pmap 下形状不匹配；编译循环掩盖 NaN；检查点缺少优化器或批统计状态。
输出内容：实验摘要、PRNG 策略、模型/状态设计、训练/评估步骤计划、最小配置、首次运行检查、JAX 特定失败模式、扩展计划。
限制：资料未提供具体代码示例、安装命令、收费或平台支持信息；实际实现需结合具体任务与框架版本自行验证。