指南# 轉換 基本範例 轉換方法 狀態傳播 轉換子狀態（提升類型） 規則和限制 軸元數據 在多個設備上擴展 概述 定義具有指定分片的模型 初始化分片模型 從檢查點載入分片模型 編譯訓練迴圈 效能分析 邏輯軸註釋 使用篩選器 篩選器協定 篩選器 DSL 群組狀態 隨機性 Rngs、RngStream 和 RngState 篩選隨機狀態 重新設定種子 分割 PRNG 金鑰 轉換 效能考量 非同步調度 降低 Python 開銷 從 Flax Linen 到 NNX 的演變 基本的 Module 定義 變數創建 訓練步驟和編譯 集合和變數類型 使用多種方法 轉換 掃描多個層 在 Flax NNX 中使用 TrainState 同時使用 Flax NNX 和 Linen 子模組就是您所需要的 基礎知識 處理 RNG 金鑰 NNX 變數類型 vs. Linen 集合 分割元數據 提升轉換 模型手術 Python 風格的 nnx.Module 操作 創建不含記憶體分配的抽象模型或狀態 檢查點手術 部分初始化 儲存和載入檢查點 設定 儲存檢查點 還原檢查點 以純字典形式儲存和還原 當檢查點結構不同時還原 多進程檢查點 其他檢查點功能 Flax NNX 與 JAX 轉換 差異 混合使用 Flax NNX 和 JAX 轉換 從 Haiku 遷移到 Flax 基本模組定義 變數創建 訓練步驟和編譯 處理非參數狀態 使用多種方法 轉換 掃描多個層 頂層 Haiku 函式 vs 頂層 Flax 模組 範例：使用預訓練的 Gemma 透過 Flax NNX 進行推論 安裝 下載模型 Python 導入 載入並準備 Gemma 模型 執行採樣/推論
