参考官方原文：Arrays and jax.numpy

先只解决一件事：把文本变成 JAX 能处理的数组，并建立对形状的敏感度。大模型代码看起来复杂，很多错误其实只是“某一维代表什么”没有想清楚。

1. JAX 是什么

JAX 是面向高性能数值计算和机器学习的 Python 库。它提供与 NumPy 相似的数组接口，同时能对我们写的函数做自动微分、自动批处理和编译，并让同一套代码运行在 CPU、GPU 或 TPU 上。

import jax
import jax.numpy as jnp

x = jnp.array([1.0, 2.0, 3.0])
print(x)
print(x.shape, x.dtype)
print(jax.devices())

[1. 2. 3.]
(3,) float32
[CpuDevice(id=0)]

如果用过 NumPy，jax.numpy 会很熟悉，但不要把 JAX 简单理解为“能跑 GPU 的 NumPy”。JAX 最重要的能力是变换整个函数。为了让函数能够被分析和变换，JAX 对程序写法有一些约束，后面会逐步解释。

传统写法是把 jax.numpy 导入为 jnp，把普通 NumPy 导入为 np：

import jax
import jax.numpy as jnp
import numpy as np

# jnp.arange(9.0) 创建一个从 0.0 到 8.0 的一维数组
# .reshape(3, 3) 将其变成一个 3 行 3 列的矩阵
x = jnp.arange(9.0).reshape(3, 3)
# 创建一个长度为 3、所有元素都是 1 的一维数组
y = jnp.ones(3)

# @ 表示矩阵乘法。这里是一个 (3, 3) 矩阵乘以一个 (3,) 向量。
print(x @ y)          # [ 3. 12. 21.]
print(x.sum(axis=0))  # [ 9. 12. 15.]
print(x[1, :2])       # [3. 4.]

JAX 数组是 jax.Array 的实例，类型标注也应使用 jax.Array。jax.typeof() 可以查看 JAX 类型；对普通数组来说，它主要表示 shape 和 dtype：

print(isinstance(x, jax.Array))  # True
print(jax.typeof(x))             # float32[3,3]

JAX 类型还能携带 sharding 信息，描述数组如何分布在设备上。NumPy 和 JAX 数组经常能通过鸭子类型互换使用：通常用 np 做数据加载与轻量 host 端处理，用 jnp 在加速器上做数值计算。

2. 数组是不可变的

NumPy 可以原地修改数组，JAX 数组默认不可变：

x = jnp.array([10, 20, 30])
# x[1] = 99  # TypeError

y = x.at[1].set(99)
print(x)  # [10 20 30]
print(y)  # [10 99 30]

x.at[1].set(99) 返回一个新值。在编译后，JAX/XLA 仍可能在底层安全地复用内存，所以“函数式写法”不等于“一定多复制一次”。

.at 除了 set，还支持 add、multiply、min 和 max：

x = jnp.arange(10)
print(x.at[3].add(100))  # [0,1,2,103,4,5,6,7,8,9]
print(x.at[:3].max(5))   # [5,5,5,3,4,5,6,7,8,9]

在 JIT 编译之外，函数式更新可能确实创建新数组；在 jax.jit() 内部，编译器通常能将更新优化为原地操作，某些中间结果甚至不会物化。确实需要可变数组时，JAX 还提供 Ref，后面的“有状态计算”会介绍。为什么这对模型重要？如果参数和中间结果不会被某段代码偷偷改掉，JAX 就更容易分析依赖关系、计算梯度并优化整个计算图。

“物化”一般是把原本抽象、延迟或尚未真正计算的东西，转换成一个实际存在、可以直接使用的结果。在 JAX 或 NumPy 语境中，它常指“真正执行计算并取得具体数组值”。

3. 从字符到 token

真正的大模型使用 BPE、SentencePiece 等 tokenizer。为了把注意力集中在 JAX 上，我们先用字符级词表：

text = "jax让模型计算更清晰"
# 提取并排序所有不同字符,得到字符表
chars = sorted(set(text))
# 建立“字符 → 编号”字典
stoi = {ch: i for i, ch in enumerate(chars)}
# 建立“编号 → 字符”字典
itos = {i: ch for ch, i in stoi.items()}

# 定义编码函数
encode = lambda s: [stoi[ch] for ch in s]
# 定义解码函数
decode = lambda ids: "".join(itos[int(i)] for i in ids)

# 把整个文本编码成 JAX 数组
tokens = jnp.array(encode(text), dtype=jnp.int32)
print(tokens)
print(decode(tokens))

[ 1  0  2 10  6  3  9  8  5  7  4]
jax让模型计算更清晰

token id 必须是整数，因为之后它要充当 embedding 表的行索引。假设词表大小为 V，隐藏维度为 D，embedding 参数的形状就是 [V, D]：

V, D = len(chars), 8
embedding_table = jnp.arange(V * D, dtype=jnp.float32).reshape(V, D)
x = embedding_table[tokens]
print(x.shape)  # [T, D]

查表以后，每个离散 token id 变成一个可学习的连续向量。

4. 构造“预测下一个 token”的样本

语言模型训练目标可以概括为：看到前面的 token，预测下一个 token。设序列长度 T=4：

data = jnp.arange(12)
T = 4

inputs = jnp.stack([data[i:i + T] for i in range(3)])
targets = jnp.stack([data[i + 1:i + T + 1] for i in range(3)])

print(inputs)
print(targets)
print(inputs.shape)  # [B, T]

inputs[b, t] 对应的监督信号是 targets[b, t]。目标序列只是输入序列向右错开一位。

5. 形状、广播与矩阵乘法

模型隐藏状态通常是 [B, T, D]，线性层权重是 [D, Dout]：

B, T, D, Dout = 2, 4, 8, 16
x = jnp.ones((B, T, D))
w = jnp.ones((D, Dout))
b = jnp.zeros((Dout,))

y = x @ w + b
print(y.shape)  # [B, T, Dout]

矩阵乘法只作用在最后一维：D 被消去。偏置 b 通过广播同时加到所有 batch 和位置上。建议在写每一行模型代码前先写形状：

# x: [B, T, D]
# w: [D, 3 * D]
qkv = x @ w  # [B, T, 3 * D]

6. dtype 不是小事

• token id 常用 int32；
• 参数和普通计算常从 float32 开始；
• 大规模训练常用 bfloat16 降低显存和计算成本，但某些归一化、损失或累加仍保留更高精度；
• Python 的整数/浮点数与 JAX 数组混合时，最终 dtype 受 JAX 的类型提升规则影响。

初学阶段不要急着上混合精度。先保证 float32 版本数值正确，再优化速度和显存。NumPy 浮点数组默认是 float64，JAX 默认是更适合加速器的 float32。JAX 默认还会禁用 64 位类型；确实需要时必须显式开启：

jax.config.update("jax_enable_x64", True)

混合 dtype 运算使用与 NumPy 相似但不完全相同的类型提升规则，特别会避免意外全部提升到 64 位：

print((jnp.arange(3) + 1.5).dtype)  # float32

7. JAX 与 NumPy 的索引差异

JAX 支持 NumPy 风格的切片和高级整数索引，但某些边缘行为不同，因为加速器内部的编译程序不能方便地抛出 Python 异常。

越界读取默认会把索引截到边界：

x = jnp.arange(10)
print(x[11])  # 9，相当于读取 x[9]

.at 的越界写入默认被丢弃：

print(x.at[11].set(99))  # x 保持不变

两种行为都可以通过索引操作的 mode 参数调整。编写训练代码时不能依赖 NumPy 的越界异常来发现数据错误，应主动检查 token id、position 和 mask 范围。

jax.numpy 也不会静默把 Python 列表转换成数组，静默转换列表容易隐藏性能问题，所以 JAX 要求显式创建数组。

# jnp.sum([1, 2, 3])  # TypeError
jnp.sum(jnp.array([1, 2, 3]))  # 正确，结果为 6

8. 数组位于设备上

JAX 数组的数据位于一个或多个 CPU、GPU 或 TPU 上。JAX 会把数组放到默认设备，计算跟随数据执行。数组也可以分片到很多设备，因此同一套全局视图代码可以从一块芯片扩展到多块芯片。

print(x.sharding)
# 单设备示例：SingleDeviceSharding(device=CpuDevice(id=0), ...)

每个数组的 sharding 属性准确描述数据放置方式。这是后续 Mesh、PartitionSpec 和 MaxText 并行配置的基础。

9. jax.numpy、jax.lax 与 XLA

jax.numpy 是面向使用者的便利层，大多数函数会处理 NumPy 风格的参数和 dtype 后，调用更严格、更底层的 jax.lax 操作。简化后的 jnp.sin 类似：

from jax import lax

def sin(x):
    x = jnp.asarray(x)
    if not jnp.issubdtype(x.dtype, jnp.inexact):
        x = x.astype(float)
    return lax.sin(x)

jax.lax 又与 XLA HLO 操作紧密对应。例如 lax.sin 基本映射到 XLA 的 Sin。因为更接近编译器，lax 不提供 jnp 的自动类型提升便利：

print(jnp.add(1, 1.0))  # 自动提升，得到 float32
# lax.add(1, 1.0)       # TypeError：两边 dtype 不同

严格性的回报是 lax 可以表达更一般的操作，包括任意 stride/dilation 的卷积和结构化控制流。多数基础代码不必直接使用它，但找不到合适的 jax.numpy API 时应检查 jax.lax。

10. 初识 JIT 编译

前面的代码可以直接运行，但此时 JAX 通常是一条一条地派发数组操作。对于会反复执行的数值计算，我们可以使用 jax.jit（Just-In-Time，即时编译）把一个 Python 函数交给 JAX 和 XLA 编译：

import jax
import jax.numpy as jnp

def linear_gelu(x, w, b):
    return jax.nn.gelu(x @ w + b)

jitted_linear_gelu = jax.jit(linear_gelu)

x = jnp.ones((2, 4))
w = jnp.ones((4, 8))
b = jnp.zeros((8,))

y = jitted_linear_gelu(x, w, b)
print(y.shape)  # (2, 8)

jax.jit 接收一个函数，返回一个功能相同但经过 JIT 编译的新函数。它不是把某一个数组“编译掉”，而是根据函数接收到的数组类型——主要包括 shape 和 dtype——追踪其中的计算，再让 XLA 生成适合当前 CPU、GPU 或 TPU 的程序。

第一次调用为什么可能更慢

JIT 函数第一次遇到某组新的 shape 和 dtype 时，需要经历追踪与编译，所以第一次调用可能比普通执行更慢。之后使用相同类型的输入，JAX 通常会复用已经编译的程序：

# 第一次：可能包含 tracing 和编译开销
y1 = jitted_linear_gelu(x, w, b)

# 第二次：shape 和 dtype 相同，通常复用编译结果
y2 = jitted_linear_gelu(x * 2, w, b)

如果把 x 从 (2, 4) 换成 (16, 4)，shape 发生变化，JAX 可能为新输入重新编译。这也是大模型训练通常保持 batch size 和 sequence length 固定，或只使用少量固定长度 bucket 的原因。

JAX 计算通常是异步派发的

在加速器上，Python 函数返回时，设备计算不一定已经完成。准确计时时需要先预热，并调用 .block_until_ready() 等待结果完成：

import time

# 预热并完成首次编译
jitted_linear_gelu(x, w, b).block_until_ready()

start = time.perf_counter()
for _ in range(100):
    y = jitted_linear_gelu(x, w, b)
y.block_until_ready()
elapsed = time.perf_counter() - start

print(f"平均耗时：{elapsed / 100:.6f} 秒")

如果不等待设备完成，只会测到 Python 把任务提交给设备所需的时间，而不是实际计算时间。

不是所有 Python 写法都能直接 JIT

JAX 追踪函数时，通常只知道数组的 shape 和 dtype，不知道运行时的具体值。因此，依赖数组具体数值的 Python if、while、int(x) 或修改全局变量等写法可能失败或产生意外行为。

后续结合 tracing、grad、vmap 和训练步骤进一步解释 JIT。这里先建立直觉：我们用熟悉的数组运算写函数，JAX 负责观察整个函数，并把它编译成适合目标设备的程序。

小结

模型不是从“Transformer”开始的，而是从数据表示开始的。先能在脑中稳定追踪 [B, T, D]，后面的注意力公式才不会变成符号堆砌。

到这里，官方“Arrays and jax.numpy”正文中的 NumPy 迁移、jax.Array/jax.typeof、不可变更新、默认 dtype、索引差异、设备放置以及 jnp—lax—XLA 层次已经全部覆盖；字符 token、embedding 和语言模型 batch 是扩展。