在瞬息萬變的機器學習與科學運算領域中，高度需要的兩個主要因素就是運算速度與彈性。來自 Google Research 團隊的 JAX 程式，在追求數值運算與深度學習領域的速度提升上，一直是名列前茅的知名工具。JAX 擁有 NumPy 的所有基本功能，除此之外，它還引入了一種非常熟悉的語言，並擁有一些最先進的功能，例如自動微分，使人們能夠運行機器學習工作負載和針對現代 GPU/TPU 優化的程式碼，同時還能夠進行即時編譯。本文將研究 JAX，以尋找最適合的使用方式，比較實際功能與聲稱的功能，並提供適合初學者的簡單教學。
什麼是 JAX？
JAX 是一個功能強大的數值運算函式庫，它擴充了 NumPy 的功能，納入了各種提升現代機器學習和科學運算生產力的函式。JAX 與舊有的函式庫不同，它是以組織函式轉移為中心，這有助於讓研究人員和工程師寫出的程式碼既簡短又強大。JAX 的主要特點是支援自動區分、硬體加速，以及容易與現有的 Python 工作流程整合。
JAX 的主要屬性
JAX 有別於其他數字計算工具包，因為它有以下獨特之處：
NumPy 相容性：JAX 支援 NumPy 所提供的大部分陣列操作。因此，使用者在使用 Numpy 運算時，不必丟棄舊有的程式碼，也可以享受 GPU/TPU 加速的好處。
自動差分：透過提供反向和正向模式的自動差分作為函式庫的一部分，JAX 降低了進入機器學習和最佳化應用程式的門檻。
即時編譯 (JIT)：jit 装饰器的功能之一是能够将 Python 函数直接编译为高效的机器代码，从而大大提高性能。
硬體加速：無需修改程式碼，JAX 就能在 CPU、GPU 和 TPU 上順暢運作，無論是簡單或大型的研究任務。
函數轉換：高效率的工具，如 grad、vmap 和 pmap，不僅可以透過向量化、平行化和自動批次化來轉換（從而優化）函數，而且可以輕鬆快速地完成轉換。

JAX 入門：基本教學
在這個例子中，我們的目標是提供一個簡短的範例，介紹 JAX 的自動區分和 JIT 編譯功能。
首先，應透過 pip 安裝 JAX：
pip install jax jaxlib
然後，應該匯入所需的模組：
匯入 jax
匯入 jax.numpy 為 jnp
從 jax 匯入 grad、jit
在尋找某些函數的梯度之前，最好先拿起一個容易理解和解釋的函數：
def f(x)：
返回 3 * x ** 2 + 2 * x + 1
df_dx = grad(f)
print(df_dx(2.0)) # 輸出：14.0
現在可以透過 JIT 編譯來賦予更快的執行速度：
@jit
def fast_f(x)：
返回 3 * x ** 2 + 2 * x + 1
print(fast_f(2.0)) # 輸出：17.0
本教學展示的不只是 JAX 在尋找漸層方面的優勢。JAX 提供 vmap 和 pmap，將這些作業融入核心，以達到批次或裝置平行化的目的。
總結
JAX 結合了 NumPy 的易用性與最新的效能最佳化，是陣列運算領域的一大進步。儘管學習如何處理函式程式設計是一個需要克服的小障礙，但其所帶來的速度與適應性程度卻相當令人滿意。無論是訓練機器學習模型、科學模擬，或是數學函數最佳化，JAX 都是有助於提升 Python 極限的工具。