JAX

JAX

開發者	Google, Nvidia
首次發布	2019年10月31日​（2019-10-31）
當前版本	0.8.1（2025年11月16日；穩定版本）
原始碼庫	github.com/jax-ml/jax
程式語言	Python, C++
作業系統	Linux, macOS, Windows
平台	Python, NumPy
類型	機器學習
許可協議	Apache 2.0
網站	docs.jax.dev/en/latest/

JAX，Python機器學習框架，用於變換數值函數，它由Google開發、還有Nvidia做出的部分貢獻。JAX結合了修改版的Autograd自動微分系統，以及來自OpenXLA專案的編譯器XLA，可加速數值線性運算。其設計目標是在介面與程式設計風格上盡可能與NumPy保持相容，使使用者能夠以熟悉的方式撰寫高效能運算程式。此外，JAX亦可與TensorFlow、PyTorch等機器學習框架整合使用。

主要功能

JAX的主要功能是：

• grad：自動微分，
• jit：即時編譯，
• vmap：自動向量化。

grad

主條目：自動微分

下面的代碼演示grad函數的自動微分。

# 导入库
from jax import grad
import jax.numpy as jnp

# 定义logistic函数
def logistic(x):  
    return jnp.exp(x) / (jnp.exp(x) + 1)

# 获得logistic函数的梯度函数
grad_logistic = grad(logistic)

# 求值logistic函数在x = 1处的梯度 
grad_log_out = grad_logistic(1.0)   
print(grad_log_out)

最終的輸出為：

0.19661194

jit

主條目：即時編譯

下面的代碼演示jit函數的優化。

# 导入库
from jax import jit
import jax.numpy as jnp

# 定义cube函数
def cube(x):
    return x * x * x

# 生成数据
x = jnp.ones((10000, 10000))

# 创建cube函数的jit版本
jit_cube = jit(cube)

# 应用cube函数和jit_cube函数于相同数据来比较其速度
cube(x)
jit_cube(x)

可見jit_cube的運行時間顯著的短於cube。

vmap

主條目：自動向量化

下面的代碼展示vmap函數的通過SIMD的向量化。

# 导入库
from functools import partial
from jax import vmap
import jax.numpy as jnp

# 定义函数
def grads(self, inputs):
    in_grad_partial = partial(self._net_grads, self._net_params)
    grad_vmap = vmap(in_grad_partial)
    rich_grads = grad_vmap(inputs)
    flat_grads = np.asarray(self._flatten_batch(rich_grads))
    assert flat_grads.ndim == 2 and flat_grads.shape[0] == inputs.shape[0]
    return flat_grads

使用JAX的庫

一些Python庫使用JAX作為後端，這包括：

• Flax，最初由Google Brain開發的高層人工神經網絡庫。
• Equinox，將參數化函數（包括人工神經網絡）表示為PyTree的庫。它由Patrick Kidger創建。
• Diffrax，用於求微分方程的數值解的庫，比如解常微分方程和隨機微分方程。
• Optax，DeepMind開發的用於梯度處理和最優化的庫。
• Lineax，用於解線性方程組和線性最小二乘法。
• RLax，DeepMind開發的用於強化學習的庫
• jraph，DeepMind開發的圖神經網絡庫。
• jaxtyping，用於為陣列或張量的形狀和數據類型增加類型標註的庫。
• NumPyro，概率編程庫。
• Brax，物理引擎。

參見

• NumPy
• TensorFlow
• PyTorch
• CUDA
• 自動微分
• 即時編譯
• 自動向量化
• 自動並行

外部連結

• Documentationː jax.readthedocs.io
• Colab (Jupyter/iPython) Quickstart Guideː colab.research.google.com/github/google/jax/blob/main/docs/notebooks/quickstart.ipynb
• TensorFlow's XLAː www.tensorflow.org/xla (Accelerated Linear Algebra)
• YouTube上的Intro to JAX: Accelerating Machine Learning research
• Original paperː mlsys.org/Conferences/doc/2018/146.pdf