第1章：JAX基础入门——从NumPy到即时编译的跨越

大家好，我是你们的老朋友。今天咱们正式开启《JAX模型部署与推理加速实战》的第一章。

说实话，我第一次接触JAX时，心里也犯嘀咕：这玩意儿跟NumPy有啥区别？不就是多了一个GPU支持吗？后来真正在项目中用起来才发现——嗯，事情没那么简单。

这一章，我会带你快速掌握JAX最核心的几个基础操作。别担心，咱们不搞花架子，全是实战干货。

1.1 jax.numpy：熟悉的配方，不一样的味道

如果你用过NumPy，那jax.numpy对你来说几乎零门槛。说白了，它就是NumPy的“JAX版本”，API几乎一模一样。

import jax.numpy as jnp
import numpy as np

# 创建数组
a = jnp.array([1, 2, 3, 4, 5])
b = jnp.ones((3, 4))
c = jnp.zeros((2, 3))

# 数学运算
d = jnp.sin(a)
e = jnp.dot(b.T, b)

print(a)
print(d)

你看，代码写起来跟NumPy完全一致。但有个关键区别——JAX的数组默认是不可变的。

注意：JAX数组不支持原地修改。比如你写 a[0] = 10 会直接报错。这是JAX函数式编程的核心设计，后面你会体会到它的好处。

我在项目中遇到过一位同事，刚转JAX时总想用 np.append 那种方式，结果调试了半天。后来我告诉他：JAX里要用 jnp.concatenate 或者 jnp.stack，而且每次操作都会返回新数组。

1.2 随机数生成与状态管理

随机数这块，JAX跟NumPy差别很大。NumPy里你写 np.random.randn() 就行，但JAX不行——它需要你显式管理随机状态。

为什么会这样？因为JAX要求所有函数都是纯函数，不能有隐藏状态。随机数生成如果依赖全局种子，那就不是纯函数了。

import jax.random as random

# 创建随机数key
key = random.PRNGKey(42)

# 生成随机数
key, subkey = random.split(key)
x = random.normal(subkey, shape=(3, 4))

print(x)
print(key)

这里有个关键点：random.split 会把一个key拆成多个，保证每个随机操作都有独立的key。我刚开始学的时候，经常忘记split，结果生成的随机数一模一样——那叫一个尴尬。

我的习惯：每次需要随机数时，先 key, subkey = random.split(key)，然后用 subkey 去生成。主key留着下次再用。这样既安全又清晰。

我曾经在一个强化学习项目里，因为随机数状态没管理好，导致训练结果无法复现。排查了两天才发现是key重复使用了。从那以后，我对随机数管理格外上心。

1.3 DeviceArray：数据在设备上的“家”

JAX里所有数组都是 DeviceArray 类型。说白了，就是数据要么在CPU上，要么在GPU/TPU上。你不需要手动 .cuda() 之类的操作，JAX会自动帮你管理。

import jax

# 查看默认设备
print(jax.devices())

# 创建数组，默认就在设备上
x = jnp.array([1, 2, 3])
print(x.device)  # 看看它在哪个设备上

# 强制放到CPU
x_cpu = jax.device_put(x, jax.devices('cpu')[0])
print(x_cpu.device)

你想想看，这比PyTorch省事多了吧？不用来回 .to(device)。JAX的哲学是：数据在哪，计算就在哪。

核心概念：DeviceArray是JAX的基石。所有操作都发生在设备上，CPU和GPU之间的数据传输是隐式的，但你可以用 jax.device_put 显式控制。

我记得第一次用JAX跑模型时，发现数据从CPU传到GPU竟然自动完成了。当时我还特意去查了文档，确认没有手动调用任何传输函数——嗯，这就是JAX的设计理念：让开发者专注于算法，而不是设备管理。

1.4 即时编译（@jit）：性能加速的开关

终于到了重头戏。@jit 是JAX最强大的功能之一。它会把你的Python函数编译成高效的XLA代码，在GPU上跑得飞快。

import jax
import jax.numpy as jnp
from jax import jit

# 定义一个普通函数
def slow_fn(x, y):
    for _ in range(1000):
        x = x * y + x
    return x

# 用jit加速
@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))

# 第一次调用会编译
result = fast_fn(x, y)

这里有个坑：第一次调用 @jit 函数时会编译，速度很慢。但后续调用会直接使用缓存，快得飞起。

避坑指南：我曾经在线上服务里用 @jit，结果第一次请求超时了——因为编译时间太长。后来我改成在服务启动时先预热一次，问题就解决了。

@jit 的原理其实不复杂：JAX会追踪函数的执行过程，生成计算图，然后交给XLA编译器优化。你想想看，这相当于白嫖了Google的编译器优化技术。

但要注意，@jit 对Python控制流有限制。比如 if 语句里的条件必须是静态的，不能依赖数组的值。我刚开始用的时候，写了个动态循环，结果编译报错——嗯，后来改成用 jax.lax.while_loop 才搞定。

1.5 本章知识体系总览

下面这张图，帮你把本章的核心知识点串起来：

JAX基础入门知识体系 JAX核心基础 jax.numpy API与NumPy一致 数组不可变 纯函数式操作 随机数生成 PRNGKey管理 random.split 可复现性 DeviceArray 自动设备管理 CPU/GPU切换 隐式传输 即时编译 @jit XLA编译优化 首次编译缓存 控制流限制 四个核心模块相互配合，构成JAX高效计算的基础

这张图把本章的四个核心模块串起来了。你仔细看，它们之间是有逻辑关系的：jax.numpy 提供基础操作，随机数管理保证可复现性，DeviceArray 负责设备管理，而 @jit 则是性能加速的关键。

1.6 本章小结

好了，第一章的内容就到这里。咱们快速回顾一下：

  • jax.numpy：跟NumPy一样用，但数组不可变
  • 随机数：必须用PRNGKey管理，记得split
  • DeviceArray：数据自动在设备上，不用手动搬运
  • @jit：一键加速，但要注意首次编译和动态控制流

这些基础操作，你在后续的模型部署和推理加速中会反复用到。尤其是 @jit，它几乎是所有性能优化的起点。

下一章，咱们会深入探讨JAX的函数式编程特性，以及如何用 vmap 和 pmap 实现自动向量化和并行化。到时候你会发现，JAX的真正威力才刚刚开始展现。

我的建议：学完这一章，最好自己动手把代码敲一遍。尤其是 @jit 那部分，感受一下编译前后的速度差异。实践出真知，这话一点不假。

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

← 上一章 📖 返回目录 下一章 →