JAX

JAX

開發者	Google, Nvidia[1]
首次发布	2019年10月31日，​5年前​（2019-10-31）[2]
当前版本	0.4.24（2024年2月6日；穩定版本）[3]
预览版本	v0.3.13（2022年5月16日，​3年前​（2022-05-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的一些贡献^[4][5][6]。它结合了修改版本的Autograd（自动通过函数的微分获得其梯度函数）^[7]，和OpenXLA的XLA（英语：Accelerated Linear Algebra）（加速线性代数）^[8]。它被设计为尽可能的遵从NumPy的结构和工作流程，并协同工作于各种现存的框架如TensorFlow和PyTorch^[9][10]。

主要功能

JAX的主要功能是^[4]：

• grad：自动微分，
• jit：即时编译，
• vmap：自动向量化，
• pmap：SPMD（英语：Single program, multiple data）编程。

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

主条目：自动向量化（英语：Automatic vectorization）

下面的代码展示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

pmap

主条目：自动并行（英语：Automatic parallelization）

下面的代码展示pmap函数的对矩阵乘法的并行化。

# 从JAX导入pmap和random；导入JAX NumPy
from jax import pmap, random
import jax.numpy as jnp

# 生成2个维度为5000 x 6000的随机数矩阵，每设备一个
random_keys = random.split(random.PRNGKey(0), 2)
matrices = pmap(lambda key: random.normal(key, (5000, 6000)))(random_keys)

# 没有数据传输，并行的在每个CPU/GPU上进行局部矩阵乘法 
outputs = pmap(lambda x: jnp.dot(x, x.T))(matrices)

# 没有数据传输，并行的在每个CPU/GPU上分别求取这两个矩阵的均值
means = pmap(jnp.mean)(outputs)
print(means)

最终的输出为：

[1.1566595 1.1805978]

使用JAX的库

一些Python库使用JAX作为后端，这包括：

• Flax，最初由Google Brain开发的高层人工神经网络库^[11]。
• Equinox，将参数化函数（包括人工神经网络）表示为PyTree的库。它由Patrick Kidger创建^[12]。
• Diffrax，用于求微分方程的数值解的库，比如解常微分方程和随机微分方程^[13]。
• Optax，DeepMind开发的用于梯度处理和最优化的库^[14]。
• Lineax，用于解线性方程组和线性最小二乘法（英语：Numerical methods for linear least squares）^[15]。
• RLax，DeepMind开发的用于强化学习的库^[16]
• jraph，DeepMind开发的图神经网络（英语：Graph neural network）库^[17]。
• jaxtyping，用于为阵列或张量的形状和数据类型增加类型标注的库^[18]。
• NumPyro，概率编程库^[19]。
• Brax，物理引擎^[20]。

参见

• NumPy
• TensorFlow
• PyTorch
• CUDA
• 自动微分
• 即时编译
• 自动向量化（英语：Automatic vectorization）
• 自动并行（英语：Automatic parallelization）

引用

1. jax/AUTHORS at main · jax-ml/jax. GitHub. [December 21, 2024]. 
2. jax-v0.1.49. 
3. https://github.com/google/jax/releases/tag/jax-v0.4.24.
4. Bradbury, James; Frostig, Roy; Hawkins, Peter; Johnson, Matthew James; Leary, Chris; MacLaurin, Dougal; Necula, George; Paszke, Adam; Vanderplas, Jake; Wanderman-Milne, Skye; Zhang, Qiao, JAX: Autograd and XLA, Astrophysics Source Code Library (Google), 2022-06-18 [2022-06-18], Bibcode:2021ascl.soft11002B, （原始内容存档于2022-06-18） 
5. Frostig, Roy; Johnson, Matthew James; Leary, Chris. Compiling machine learning programs via high-level tracing (PDF). MLsys. 2018-02-02: 1–3. （原始内容存档 (PDF)于2022-06-21）.  引文格式1维护：日期与年 (link)
6. Using JAX to accelerate our research. www.deepmind.com. [2022-06-18]. （原始内容存档于2022-06-18） （英语）. 
7. autograd. [2023-09-23]. （原始内容存档于2022-07-18）. 
8. XLA. [2023-09-23]. （原始内容存档于2022-09-01）. 
9. Lynley, Matthew. Google is quietly replacing the backbone of its AI product strategy after its last big push for dominance got overshadowed by Meta. Business Insider. [2022-06-21]. （原始内容存档于2022-06-21） （美国英语）. 
10. Why is Google's JAX so popular?. Analytics India Magazine. 2022-04-25 [2022-06-18]. （原始内容存档于2022-06-18） （美国英语）. 
11. Flax: A neural network library and ecosystem for JAX designed for flexibility, Google, 2022-07-29 [2022-07-29], （原始内容存档于2022-09-03） 
12. Kidger, Patrick, Equinox, 2022-07-29 [2022-07-29], （原始内容存档于2023-09-19） 
13. Kidger, Patrick, Diffrax, 2023-08-05 [2023-08-08], （原始内容存档于2023-08-10） 
14. Optax, DeepMind, 2022-07-28 [2022-07-29], （原始内容存档于2023-06-07） 
15. Lineax, Google, 2023-08-08 [2023-08-08], （原始内容存档于2023-08-10） 
16. RLax, DeepMind, 2022-07-29 [2022-07-29], （原始内容存档于2023-04-26） 
17. Jraph - A library for graph neural networks in jax., DeepMind, 2023-08-08 [2023-08-08], （原始内容存档于2022-11-23） 
18. jaxtyping, Google, 2023-08-08 [2023-08-08], （原始内容存档于2023-08-10） 
19. NumPyro － Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU. [2022-08-31]. （原始内容存档于2022-08-31）. 
20. Brax － Massively parallel rigidbody physics simulation on accelerator hardware. [2022-08-31]. （原始内容存档于2022-08-31）. 

外部链接

• 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