指南# 转换 基本示例 转换方法 状态传播 转换子状态（提升类型） 一致别名 轴元数据 在多个设备上扩展 概述 定义具有指定分片的模型 初始化分片模型 从检查点加载分片模型 编译训练循环 性能分析 逻辑轴注释 使用过滤器 过滤器协议 过滤器 DSL 分组状态 随机性 Rngs、RngStream 和 RngState 过滤随机状态 重新播种 拆分 Rngs 转换 从 Linen 到 NNX 的演变 基本模块定义 变量创建 训练步骤和编译 集合和变量类型 使用多种方法 提升转换 跨层扫描 在 NNX 中使用 TrainState 一起使用 Flax NNX 和 Linen 子模块是你需要的全部 基础知识 处理 RNG 密钥 NNX 变量类型与 Linen 集合 分区元数据 提升转换 模型手术 Pythonic 模块操作 创建没有内存分配的抽象模型或状态 检查点手术 部分初始化 Flax NNX vs JAX 转换 NNX 和 JAX 转换之间的区别 混合使用 Flax NNX 和 JAX 转换 从 Haiku 迁移到 Flax 基本示例 处理状态 使用多种方法 提升转换 跨层扫描 顶级 Haiku 函数与顶级 Flax 模块
