第2章 JAX核心基石

第2章：JAX核心基石：NumPy兼容API、设备数组、即时编译(JIT)入门

好，咱们正式开始动手了。上一章我们聊了JAX的定位和生态，这一章，我们直接切入它的三个核心能力。说白了，这三个东西就是JAX的「三驾马车」：长得像NumPy的API、能跑在GPU上的数组、以及那个神奇的即时编译。

我个人习惯，学任何框架都先看它的数据结构和基本运算。JAX这点做得特别聪明——它让你几乎不用重新学习。你想想看，如果你已经会NumPy，那JAX的入门成本几乎为零。

2.1 NumPy兼容API：熟悉的配方，不一样的味道

JAX的API设计哲学很直接：能复用NumPy的，绝不另起炉灶。它提供了一个叫 jax.numpy 的模块，别名通常叫 jnp。你之前怎么写NumPy，现在就怎么写JAX。

  核心要点： jax.numpy 覆盖了NumPy 90%以上的常用函数，包括数组创建、数学运算、线性代数、随机数等。

来看个例子，感受一下这种「无缝切换」的感觉：

import jax.numpy as jnp
import numpy as np

# 创建数组，和NumPy一模一样
a = jnp.array([1.0, 2.0, 3.0])
b = jnp.array([4.0, 5.0, 6.0])

# 基本运算，完全一致
c = a + b
d = jnp.dot(a, b)
e = jnp.sum(a)

print(c)  # [5. 7. 9.]
print(d)  # 32.0

嗯，这里要注意：虽然长得像，但底层完全不同。NumPy的数组是存在CPU内存里的，而JAX的数组默认是存在设备（GPU/TPU）上的。这个区别，引出了我们下一个核心概念。

2.2 设备数组：数据在哪儿，计算就在哪儿

JAX里的数组，官方叫 DeviceArray。什么意思呢？就是数组本身绑定了一个设备。你在GPU上创建的数组，计算默认就在GPU上跑。

我在项目中遇到过一个问题：刚开始用JAX时，我习惯性地用 np.array 创建数据，然后传给JAX函数。结果发现每次都要从CPU拷贝到GPU，慢得离谱。后来才意识到，应该直接用 jnp.array 创建。

我的建议： 在JAX项目中，尽量全程使用 jnp 操作数据。混用 np 和 jnp 会导致隐式的设备传输，拖慢性能。

你可以显式地查看数组在哪个设备上：

import jax

x = jnp.array([1, 2, 3])
print(x.device())  # 比如：cuda:0 或 tpu:0

# 你也可以手动指定设备
with jax.default_device(jax.devices("cpu")[0]):
    y = jnp.array([4, 5, 6])
    print(y.device())  # cpu:0

为什么会这样设计？说白了，就是为了让计算靠近数据。你想想看，在大规模模型训练时，数据在GPU显存里，计算也在GPU上，省去了来回搬运的开销。这就是JAX高性能的根基之一。

2.3 即时编译(JIT)：让Python代码飞起来

好了，前两个概念还算温和。接下来这个，才是JAX真正让人兴奋的地方——即时编译。

JAX的 jit 装饰器，可以把你的Python函数编译成高效的XLA（Accelerated Linear Algebra）计算图。第一次调用时编译，后续调用直接执行编译后的代码，速度能提升几十倍甚至上百倍。

我曾经在做一个序列标注模型时，一个前向传播函数用纯Python写要跑200ms。加上 @jit 之后，直接降到5ms。当时我盯着屏幕愣了好几秒——这玩意儿是真的猛。

  避坑指南： jit 编译的函数中，不能有Python控制流（if/for/while）依赖动态数据。因为编译时需要静态的计算图。如果必须用动态控制流，请使用 jax.lax.cond 或 jax.lax.scan 等函数。

来看一个完整的JIT示例：

import jax
import jax.numpy as jnp
import time

# 定义一个简单的函数
def slow_fn(x, y):
    # 模拟一些计算
    for _ in range(1000):
        x = x * y + x
    return x

# 用jit装饰
@jax.jit
def fast_fn(x, y):
    for _ in range(1000):
        x = x * y + x
    return x

# 准备数据
x = jnp.ones((1000, 1000))
y = jnp.ones((1000, 1000))

# 测试慢版本
start = time.time()
result_slow = slow_fn(x, y)
print(f"未编译耗时: {time.time() - start:.3f}秒")

# 第一次调用会编译
start = time.time()
result_fast = fast_fn(x, y)
print(f"第一次JIT调用(含编译): {time.time() - start:.3f}秒")

# 后续调用直接执行编译后的代码
start = time.time()
result_fast = fast_fn(x, y)
print(f"第二次JIT调用(纯执行): {time.time() - start:.3f}秒")

你运行一下就会发现，第二次JIT调用的速度，比未编译版本快了一个数量级。这就是JIT的威力。

2.4 三者的关系：一张图说清楚

这三个概念不是孤立的，它们环环相扣。我画了一张图，帮你理清它们的关系：

这张图很直观：你用熟悉的NumPy API创建数据，数据天然就在设备上，然后通过JIT编译让计算飞起来。三者缺一不可。

2.5 实战小练习：验证你的理解

光看不练假把式。我建议你立刻打开终端，跑一下这段代码，感受一下JIT带来的性能差异：

import jax
import jax.numpy as jnp
import time

# 定义一个矩阵乘法函数
def matmul_fn(a, b):
    return jnp.dot(a, b)

# JIT版本
@jax.jit
def matmul_fn_jit(a, b):
    return jnp.dot(a, b)

# 创建大矩阵
a = jnp.ones((2000, 2000))
b = jnp.ones((2000, 2000))

# 预热（JIT第一次调用会编译）
_ = matmul_fn_jit(a, b)

# 测试非JIT版本
start = time.time()
for _ in range(100):
    _ = matmul_fn(a, b)
print(f"非JIT版本: {(time.time() - start)/100*1000:.3f} ms/次")

# 测试JIT版本
start = time.time()
for _ in range(100):
    _ = matmul_fn_jit(a, b)
print(f"JIT版本: {(time.time() - start)/100*1000:.3f} ms/次")

你猜结果会差多少倍？我当初第一次跑的时候，JIT版本快了将近50倍。这就是为什么我说JAX是「为性能而生」的框架。

注意： JIT编译后的函数，输入数据的形状和类型必须保持一致。如果你传入不同形状的数组，JAX会重新编译，导致性能下降。所以，尽量保持输入数据形状稳定。

好了，这一章的内容就到这里。你掌握了JAX的三个核心概念：NumPy兼容API让你快速上手，设备数组让数据靠近计算，JIT让代码飞起来。下一章，我们会深入JIT的更多细节，包括如何调试、如何避免常见的编译陷阱。

公众号:蓝海资料掘金营，微信deep3321

📚 课程目录