JAX，由Google开发、並由Nvidia做出部分貢獻[4][5][6]的Python机器学习框架，用於变换数值函数。JAX结合了修改版的Autograd自動微分系統[7]，以及來自OpenXLA專案的編譯器XLA[8]，可加速數值線性運算。其設計目標是在介面與程式設計風格上盡可能與NumPy保持相容，使使用者能夠以熟悉的方式撰寫高效能運算程式。此外，JAX亦可與TensorFlow、PyTorch等機器學習框架整合使用。[9][10]
 | |
| 開發者 | Google, Nvidia[1] |
|---|---|
| 首次发布 | 2019年10月31日[2] |
| 当前版本 |
|
| 源代码库 | github |
| 编程语言 | Python, C++ |
| 操作系统 | Linux, macOS, Windows |
| 平台 | Python, NumPy |
| 类型 | 机器学习 |
| 许可协议 | Apache 2.0 |
| 网站 | docs |
主要功能
编辑JAX的主要功能是[4]：
- grad：自动微分，
- jit：即时编译，
- vmap：自动向量化，
- pmap：SPMD编程。
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开发的高层人工神经网络库[11]。
- Equinox，将参数化函数（包括人工神经网络）表示为PyTree的库。它由Patrick Kidger创建[12]。
- Diffrax，用于求微分方程的数值解的库，比如解常微分方程和随机微分方程[13]。
- Optax，DeepMind开发的用于梯度处理和最优化的库[14]。
- Lineax，用于解线性方程组和线性最小二乘法[15]。
- RLax，DeepMind开发的用于强化学习的库[16]
- jraph，DeepMind开发的图神经网络库[17]。
- jaxtyping，用于为阵列或张量的形状和数据类型增加类型标注的库[18]。
- NumPyro，概率编程库[19]。
- Brax，物理引擎[20]。
参见
编辑引用
编辑- ^ jax/AUTHORS at main · jax-ml/jax. GitHub. [December 21, 2024].
- ^ jax-v0.1.49.
- ^ https://github.com/jax-ml/jax/releases/tag/jax-v0.8.1.
- ^ 4.0 4.1 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）
- ^ 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）.
- ^ Using JAX to accelerate our research. www.deepmind.com. [2022-06-18]. （原始内容存档于2022-06-18） （英语）.
- ^ autograd. [2023-09-23]. （原始内容存档于2022-07-18）.
- ^ XLA. [2023-09-23]. （原始内容存档于2022-09-01）.
- ^ 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） （美国英语）.
- ^ Why is Google's JAX so popular?. Analytics India Magazine. 2022-04-25 [2022-06-18]. （原始内容存档于2022-06-18） （美国英语）.
- ^ Flax: A neural network library and ecosystem for JAX designed for flexibility, Google, 2022-07-29 [2022-07-29], （原始内容存档于2022-09-03）
- ^ Kidger, Patrick, Equinox, 2022-07-29 [2022-07-29], （原始内容存档于2023-09-19）
- ^ Kidger, Patrick, Diffrax, 2023-08-05 [2023-08-08], （原始内容存档于2023-08-10）
- ^ Optax, DeepMind, 2022-07-28 [2022-07-29], （原始内容存档于2023-06-07）
- ^ Lineax, Google, 2023-08-08 [2023-08-08], （原始内容存档于2023-08-10）
- ^ RLax, DeepMind, 2022-07-29 [2022-07-29], （原始内容存档于2023-04-26）
- ^ Jraph - A library for graph neural networks in jax., DeepMind, 2023-08-08 [2023-08-08], （原始内容存档于2022-11-23）
- ^ jaxtyping, Google, 2023-08-08 [2023-08-08], （原始内容存档于2023-08-10）
- ^ NumPyro － Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU. [2022-08-31]. （原始内容存档于2022-08-31）.
- ^ 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
