開始使用 JAX

歡迎使用 JAX！JAX 文件包含許多有用的資源，可協助您開始使用。快速入門是最容易上手並概覽 JAX 專案的地方。

如果您習慣編寫 NumPy 程式碼，並且剛開始探索 JAX，您可能會發現以下資源很有幫助

• 如何以 JAX 思考是對 JAX 執行模型的概念性演練。

• 🔪 JAX - 尖銳之處 🔪 列出了一些 JAX 的尖銳角落。

• 常見問題 (FAQ) 回答了一些常見的 jax 問題。

教學

如果您準備好更深入地探索 JAX，JAX 教學會深入探討更多細節

• 教學
  • 快速入門
  • 重要概念
  • 即時編譯
  • 自動向量化
  • 自動微分
  • 偵錯簡介
  • 虛擬隨機數
  • 使用 pytrees
  • 平行程式設計簡介
  • 具狀態的計算
  • 使用 JIT 的控制流程和邏輯運算子
  • 進階自動微分
  • 外部回呼
  • 使用 jax.checkpoint (jax.remat) 的梯度檢查點
  • JAX 內部機制：primitives
  • JAX 內部機制：jaxpr 語言

如果您喜歡影片介紹，這裡有 JAX 貢獻者 Jake VanderPlas 的一個影片

基於 JAX 建構

JAX 為更大的社群開發的許多工具提供了核心數值運算 primitives。例如，如果您有興趣使用 JAX 訓練神經網路，兩個受到良好支援的選項是 Flax 和 Haiku。

如需社群策劃的跨多個領域的 JAX 相關專案列表，請查看 Awesome JAX。

尋求協助

如果您對 JAX 有疑問，我們很樂意回答！以下是兩個可以獲得解答的好地方

• JAX GitHub 討論區

• StackOverflow 上的 JAX