進階指南

本節包含關於更進階主題的範例和教學，例如多核心計算、自動微分和自訂運算。

平行計算

• 分散式陣列和自動平行化
• 使用 shard_map 的手動平行化
• 多主機和多進程環境
• 分散式資料載入

自動微分

• 自動微分食譜
• 自訂導數規則
• 使用 jax.checkpoint (又名 jax.remat) 控制自動微分的已儲存值

深入探討

• JAX 中的廣義卷積
• XLA 編譯器標誌列表