快速入门#
JAX 是一个面向数组的数值计算库（类似于 NumPy），具有自动微分和 JIT 编译功能，可实现高性能的机器学习研究.
本文档提供 JAX 基本功能的快速概述，以便您可以快速开始使用 JAX
JAX 为在 CPU、GPU 或 TPU 上运行的计算提供统一的类似 NumPy 的接口，无论是在本地还是分布式环境中。
JAX 通过 Open XLA（一个开源机器学习编译器生态系统）提供内置的即时 (JIT) 编译。
JAX 函数通过其自动微分转换支持高效的梯度评估。
JAX 函数可以自动向量化，以有效地将它们映射到表示一批输入的数组上。
安装#
可以直接从 Python 包索引安装 JAX 的 CPU 版本，适用于 Linux、Windows 和 macOS
pip install jax
或者，对于 NVIDIA GPU
pip install -U "jax[cuda12]"
有关更详细的特定于平台的安装信息，请查看 安装。
JAX 作为 NumPy#
大多数 JAX 的使用是通过熟悉的 jax.numpy API 进行的，它通常以 jnp 别名导入。
import jax.numpy as jnp
通过此导入，您可以立即以类似于典型 NumPy 程序的方式使用 JAX，包括使用 NumPy 风格的数组创建函数、Python 函数和运算符，以及数组属性和方法。
def selu(x, alpha=1.67, lmbda=1.05):
return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)
x = jnp.arange(5.0)
print(selu(x))
[0. 1.05 2.1 3.1499999 4.2 ]
一旦您开始深入研究，您会发现 JAX 数组和 NumPy 数组之间存在一些差异；这些差异在 🔪 JAX - 锋利之处 🔪 中进行了探讨。
使用 jax.jit() 进行即时编译#
JAX 在 GPU 或 TPU 上透明地运行（如果没有，则回退到 CPU）。但是，在上面的示例中，JAX 正在一次一个操作地将内核分派到芯片。如果有一系列操作，我们可以使用 jax.jit() 函数，使用 XLA 将这一系列操作一起编译。
我们可以使用 IPython 的 %timeit 来快速基准测试我们的 selu 函数，使用 block_until_ready() 来考虑 JAX 的动态分派（请参阅 异步分派）
from jax import random
key = random.key(1701)
x = random.normal(key, (1_000_000,))
%timeit selu(x).block_until_ready()
2.54 ms ± 49.5 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
（注意，我们使用了 jax.random 来生成一些随机数；有关如何在 JAX 中生成随机数的详细信息，请查看 伪随机数）。
我们可以使用 jax.jit() 转换来加速此函数的执行，它会在第一次调用 selu 时进行即时编译，并在之后进行缓存。
from jax import jit
selu_jit = jit(selu)
_ = selu_jit(x) # compiles on first call
%timeit selu_jit(x).block_until_ready()
569 μs ± 4.36 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
上面的计时代表在 CPU 上执行，但相同的代码可以在 GPU 或 TPU 上运行，通常会有更大的加速。
有关 JAX 中 JIT 编译的更多信息，请查看 即时编译。
使用 jax.grad() 求导数#
除了通过 JIT 编译转换函数外，JAX 还提供了其他转换。其中一种转换是 jax.grad()，它执行自动微分（autodiff）
from jax import grad
def sum_logistic(x):
return jnp.sum(1.0 / (1.0 + jnp.exp(-x)))
x_small = jnp.arange(3.)
derivative_fn = grad(sum_logistic)
print(derivative_fn(x_small))
[0.25 0.19661197 0.10499357]
让我们通过有限差分来验证我们的结果是否正确。
def first_finite_differences(f, x, eps=1E-3):
return jnp.array([(f(x + eps * v) - f(x - eps * v)) / (2 * eps)
for v in jnp.eye(len(x))])
print(first_finite_differences(sum_logistic, x_small))
[0.24998187 0.1965761 0.10502338]
grad() 和 jit() 转换可以组合，并且可以任意混合。在上面的示例中，我们对 sum_logistic 进行了 jit 编译，然后求了它的导数。我们可以更进一步。
print(grad(jit(grad(jit(grad(sum_logistic)))))(1.0))
-0.0353256
除了标量值函数之外，jax.jacobian() 转换可用于计算向量值函数的完整雅可比矩阵。
from jax import jacobian
print(jacobian(jnp.exp)(x_small))
[[1. 0. 0. ]
[0. 2.7182817 0. ]
[0. 0. 7.389056 ]]
对于更高级的自动微分操作，您可以使用 jax.vjp() 进行反向模式向量雅可比积，以及使用 jax.jvp() 和 jax.linearize() 进行前向模式雅可比向量积。两者可以彼此任意组合，并且可以与其他 JAX 转换组合。例如，jax.jvp() 和 jax.vjp() 用于定义前向模式 jax.jacfwd() 和反向模式 jax.jacrev()，分别用于以正向和反向模式计算雅可比矩阵。以下是一种组合它们以创建一个可以有效地计算完整 Hessian 矩阵的函数的方法。
from jax import jacfwd, jacrev
def hessian(fun):
return jit(jacfwd(jacrev(fun)))
print(hessian(sum_logistic)(x_small))
[[-0. -0. -0. ]
[-0. -0.09085776 -0. ]
[-0. -0. -0.07996249]]
这种组合在实践中会产生高效的代码；这或多或少是 JAX 内置的 jax.hessian() 函数的实现方式。
有关 JAX 中自动微分的更多信息，请查看 自动微分。
使用 jax.vmap() 进行自动向量化#
另一个有用的转换是 vmap()，即向量化映射。它具有沿数组轴映射函数的熟悉语义，但它不是显式地循环函数调用，而是将函数转换为本地向量化版本以获得更好的性能。当与 jit() 组合时，它可以像手动重写函数以对额外的批次维度进行操作一样高效。
我们将使用一个简单的示例，并使用 vmap() 将矩阵向量积提升为矩阵矩阵积。尽管在这种情况中很容易手动完成，但相同的技术可以应用于更复杂的函数。
key1, key2 = random.split(key)
mat = random.normal(key1, (150, 100))
batched_x = random.normal(key2, (10, 100))
def apply_matrix(x):
return jnp.dot(mat, x)
apply_matrix 函数将向量映射到向量，但我们可能希望将其按行应用于矩阵。我们可以通过在 Python 中循环批次维度来实现这一点，但这通常会导致性能不佳。
def naively_batched_apply_matrix(v_batched):
return jnp.stack([apply_matrix(v) for v in v_batched])
print('Naively batched')
%timeit naively_batched_apply_matrix(batched_x).block_until_ready()
Naively batched
636 μs ± 6.86 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
熟悉 jnp.dot 函数的程序员可能会意识到，可以使用 jnp.dot 的内置批处理语义重写 apply_matrix，从而避免显式循环。
import numpy as np
@jit
def batched_apply_matrix(batched_x):
return jnp.dot(batched_x, mat.T)
np.testing.assert_allclose(naively_batched_apply_matrix(batched_x),
batched_apply_matrix(batched_x), atol=1E-4, rtol=1E-4)
print('Manually batched')
%timeit batched_apply_matrix(batched_x).block_until_ready()
Manually batched
28.7 μs ± 2.15 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
但是，随着函数变得越来越复杂，这种手动批处理变得越来越困难且容易出错。vmap() 转换旨在自动将函数转换为批处理感知版本。
from jax import vmap
@jit
def vmap_batched_apply_matrix(batched_x):
return vmap(apply_matrix)(batched_x)
np.testing.assert_allclose(naively_batched_apply_matrix(batched_x),
vmap_batched_apply_matrix(batched_x), atol=1E-4, rtol=1E-4)
print('Auto-vectorized with vmap')
%timeit vmap_batched_apply_matrix(batched_x).block_until_ready()
Auto-vectorized with vmap
37.9 μs ± 1.02 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
正如您所期望的那样，vmap() 可以与 jit()、grad() 和任何其他 JAX 转换进行任意组合。
有关 JAX 中自动向量化的更多信息，请查看 自动向量化。
这只是 JAX 功能的冰山一角。我们真的很期待看到您如何使用它！