關於專案

關於專案#

JAX 專案由 JAX 核心團隊領導。我們以開放方式開發，並歡迎來自社群各界的開放原始碼貢獻。我們經常看到來自 Google DeepMind、更廣泛的 Alphabet、NVIDIA 和其他地方的貢獻。

專案的核心是 JAX 核心函式庫，其重點在於機器學習和數值計算的基礎，並具有規模。

在開發核心時，我們希望保持敏捷性和重點範圍，因此我們非常依賴周圍的模組化技術堆疊。首先，我們將 jax 模組設計為可組合和可擴展，以便各種特定領域的函式庫能夠以分散式的方式在其外部蓬勃發展。其次，我們非常依賴模組化後端堆疊（編譯器和執行階段）來鎖定不同的加速器。無論您是撰寫使用 JAX 建構的新的特定領域函式庫，還是希望支援新的硬體，您通常可以在幾乎不修改 JAX 核心程式碼庫的情況下貢獻這些內容。

JAX 的許多核心貢獻者都紮根於開放原始碼軟體和研究，領域涵蓋電腦科學和自然科學。我們致力於持續推動機器學習和數值計算的前沿——跨所有運算平台和加速器——並發現大規模陣列程式設計的真諦。

開放開發#

JAX 的日常開發在 GitHub 上公開進行，使用 pull request、issue 追蹤器、討論和JAX 增強提案 (JEPs)。閱讀和參與這些是參與的好方法。我們也維護了涵蓋 JAX 內部設計的開發者筆記。

JAX 核心團隊決定是否接受變更和增強功能。維持簡單的決策結構目前有助於我們以研究前沿的速度進行開發。開放開發是我們的核心價值，如果/當更複雜的決策結構（例如，指定區域負責人）變得有用時，我們可能會適應它。

如需更多資訊，請參閱貢獻 JAX。

模組化堆疊#

為了實現 (a) 跨數值領域不斷成長的使用者社群，以及 (b) 不斷發展的硬體環境，我們非常依賴模組化。

基於 JAX 建構的函式庫#

雖然 JAX 核心函式庫專注於基礎，但我們希望鼓勵在 JAX 之上建構特定領域的函式庫和工具。實際上，許多函式庫已經圍繞 JAX 出現，以提供更高等級的功能和擴充。

我們如何鼓勵這種分散式開發？我們透過幾個技術選擇來引導它。首先，JAX 的主要 API 專注於基本建構區塊（例如，數值 primitives、NumPy 運算、陣列和轉換），鼓勵輔助函式庫根據其領域的需求開發實用程式。此外，JAX 公開了一些更進階的 API，用於自訂和擴展性。函式庫可以依賴這些 API，以便將 JAX 用作內部實作手段，以更深入地與其轉換（如自動微分）整合等等。

JAX 生態系統中的專案以分散式且通常開放的方式開發。它們不受 JAX 核心團隊的管轄，即使有時團隊成員會為其做出貢獻或與其開發人員保持聯繫。

可插拔的後端#

我們希望 JAX 能夠在 CPU、GPU、TPU 和其他新興硬體平台上執行。為了鼓勵在新平台上不受阻礙地支援 JAX，JAX 核心也強調其後端的模組化。

為了管理硬體裝置和記憶體，以及針對這些裝置進行編譯，JAX 呼叫開放的 XLA 編譯器和 PJRT 執行階段。這兩者都是 JAX 核心之外的專案，由 OpenXLA 管轄和維護（同樣，JAX 核心開發人員經常做出貢獻並與之討論）。

XLA 旨在實現跨加速器的互通性（例如，透過攝取 StableHLO 作為輸入），而 PJRT 透過外掛程式裝置 API 提供擴展性。新增對新裝置的支援是透過實作 XLA 的後端降低，以及實作 PJRT 定義的外掛程式裝置 API 來完成的。如果您希望為編譯做出貢獻，或支援新的硬體，我們鼓勵您在 XLA 和 PJRT 層面做出貢獻。

這些開放系統組件允許第三方在新加速器平台上支援 JAX，而無需變更 JAX 核心。今天有幾個外掛程式正在開發中。例如，Apple 的一個團隊正在開發 PJRT 外掛程式，以使 JAX 在 Apple Metal 上執行。