快速入门

快速入门#

JAX 是一个面向数组的数值计算库（类似于 NumPy），具有自动微分和即时编译功能，以支持高性能机器学习研究。

本文档提供了JAX关键特性的快速概览，以便您能快速上手JAX：

  • JAX 提供了一个统一的类似 NumPy 的接口，用于在 CPU、GPU 或 TPU 上运行的计算，无论是在本地还是分布式环境中。

  • JAX 通过 Open XLA 提供了内置的即时（JIT）编译功能，这是一个开源的机器学习编译器生态系统。

  • JAX 函数通过其自动微分变换支持高效的梯度评估。

  • JAX 函数可以自动向量化，以高效地映射代表输入批次的数组。

安装#

JAX 可以在 Linux、Windows 和 macOS 上为 CPU 直接从 Python 包索引 安装：

pip install jax

或者，对于 NVIDIA GPU：

pip install -U "jax[cuda12]"

有关更详细的平台特定安装信息，请查看 安装。

JAX 作为 NumPy#

大多数 JAX 的使用是通过熟悉的 jax.numpy API 进行的，该 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 - The Sharp Bits 🔪 中进行了探讨。

使用 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()

1.31 ms ± 104 μs per loop (mean ± std. dev. of 7 runs, 1,000 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()

484 μs ± 124 μ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.1964569  0.10502338]

The 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.09085775 -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
393 μs ± 18.2 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

熟悉 jnp.dot 函数的程序员可能会意识到，apply_matrix 可以通过使用 jnp.dot 的内置批处理语义来重写，以避免显式循环：

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
22.9 μs ± 2.83 μ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
26.3 μs ± 2.06 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

正如你所期望的，vmap() 可以与 jit()、grad() 以及任何其他 JAX 变换任意组合。

有关JAX中自动矢量化的更多信息，请查看 自动矢量化。

这只是 JAX 能做的一小部分。我们真的很期待看到你用它做什么！

目录