3. 核心转换机制（上）：jit即时编译的原理、jit的使用与性能分析、jit的静态参数与动态形状

好，咱们今天聊聊 JAX 里最核心、也最让人又爱又恨的东西——jit。

说白了，jit 就是 JAX 的“性能引擎”。没有它，JAX 就是个普通的 NumPy 替代品。有了它，你才能把代码跑在 GPU、TPU 上，甚至还能自动做算子融合、内存优化。我刚开始接触 JAX 时，总觉得 jit 就是个装饰器，套上就完事了。后来踩了不少坑才明白——jit 的底层逻辑，决定了你代码的上限。

3.1 jit 即时编译的原理

jit 的全称是 Just-In-Time compilation。但 JAX 的 jit 和传统意义上的 JIT（比如 Java 的 JIT）不太一样。它其实是 “先追踪，再编译，最后缓存” 的过程。

具体来说，当你用 @jit 装饰一个函数，JAX 并不会立刻执行它。而是先通过 JAXPR（JAX 的内部中间表示） 把你的 Python 函数“拍平”成一个计算图。这个图是纯函数式的，没有副作用，没有控制流（除非你用 lax.cond 或 lax.scan）。

然后，XLA（Accelerated Linear Algebra）编译器接手。XLA 会把 JAXPR 图做一系列优化：

• 算子融合：把多个小算子合并成一个大的 kernel，减少显存带宽消耗。
• 内存规划：提前分配好 buffer，避免动态分配的开销。
• 设备代码生成：针对 GPU/TPU 生成对应的 PTX 或 HLO 代码。

最后，编译好的可执行文件会被缓存起来。下次调用同样的函数、同样的参数形状时，直接命中缓存，跳过编译。

我在项目中遇到过一个问题：一个模型在训练时每步都重新编译一次，导致训练速度慢得离谱。后来发现是因为输入数据的 batch size 偶尔会变化，触发了重新编译。嗯，这个坑咱们后面会细讲。

3.2 jit 的使用与性能分析

使用 jit 非常简单。最常见的两种方式：

import jax
import jax.numpy as jnp

# 方式一：装饰器
@jax.jit
def f(x, y):
    return jnp.dot(x, y) + x

# 方式二：函数式调用
f_jitted = jax.jit(f)

我个人习惯用装饰器，代码更干净。但如果你需要动态调整 jit 参数（比如 static_argnums），函数式调用会更灵活。

但光会用还不够。你得学会分析 jit 的性能。JAX 提供了几个好用的工具：

• jax.jit(f).lower(...).compile()：手动触发编译，查看编译时间。
• jax.xla_computation(f)：导出 HLO 图，看看 XLA 到底做了什么优化。
• jax.profiler：配合 Nsight 或 TensorBoard，分析 kernel 执行时间。

举个例子，你想知道 jit 编译花了多久：

import time

@jax.jit
def g(x):
    return jnp.sin(x) + jnp.cos(x)

x = jnp.ones((1000, 1000))

start = time.time()
_ = g(x)  # 第一次调用，触发编译
print(f"第一次调用耗时: {time.time() - start:.4f} 秒")

start = time.time()
_ = g(x)  # 第二次调用，命中缓存
print(f"第二次调用耗时: {time.time() - start:.4f} 秒")

你会发现，第一次调用可能花了几百毫秒甚至几秒（取决于模型复杂度），第二次调用就变成微秒级了。

我曾经在优化一个 Transformer 推理管线时，发现每次请求都要重新编译。用 profiler 一查，原来是输入序列长度不固定，导致 XLA 每次都生成不同的计算图。后来我用了 padding + 动态形状 的技巧，才把编译次数降下来。

3.3 jit 的静态参数与动态形状

这是 jit 里最容易踩坑的地方，也是面试常考题。咱们好好捋一捋。

静态参数：指的是那些在编译时就必须确定的值。比如数组的形状、循环的次数、条件判断的分支。JAX 要求这些值在编译时是已知的、固定的。

动态形状：指的是运行时才能确定的形状。比如一个 batch 的数据，有的样本长、有的短。如果你直接传给 jit 函数，JAX 会认为形状变了，触发重新编译。

为什么会这样？因为 XLA 编译器需要知道每个 tensor 的形状，才能分配内存、规划计算。形状一变，整个计算图都得重新生成。

那怎么处理动态形状呢？有几种常见策略：

1. Padding（填充）：把所有样本填充到最大长度。这是最常用的方法。缺点是浪费计算和显存。
2. Bucketing（分桶）：把相似长度的样本分到同一个 bucket，每个 bucket 内 padding 到该 bucket 的最大长度。这样编译次数可控（最多 bucket 数量次），浪费也少。
3. 使用 static_argnums：把形状参数标记为静态。这样你可以把形状作为参数传进去，但每次形状变化都会触发重新编译。

看个例子：

@jax.jit
def h(x, y):
    return x + y

# 动态形状：每次调用形状不同，会重新编译
h(jnp.ones((3,)), jnp.ones((3,)))
h(jnp.ones((4,)), jnp.ones((4,)))  # 重新编译！

# 使用 static_argnums 标记形状参数
@partial(jax.jit, static_argnums=(2,))
def h_static(x, y, shape):
    x = x.reshape(shape)
    y = y.reshape(shape)
    return x + y

h_static(jnp.ones(12), jnp.ones(12), (3, 4))
h_static(jnp.ones(12), jnp.ones(12), (4, 3))  # 重新编译，因为 shape 变了

我个人建议：能用 padding 就用 padding。实在不行再考虑 bucketing。至于 static_argnums，只适合那些很小的、不会频繁变化的参数（比如配置项、布尔标志）。

最后，咱们用一张图总结一下 jit 的核心流程：

你看，整个流程其实不复杂。但真正用好 jit，需要你对静态 vs 动态有深刻理解。记住一句话：jit 喜欢确定性，讨厌变化。你的代码越“静态”，jit 就越高效。

嗯，这一节咱们把 jit 的原理、使用、性能分析和静态参数都过了一遍。下一节咱们会深入 vmap 和 pmap，看看 JAX 怎么把单机计算扩展到多设备。到时候你会发现，jit 和 vmap 配合起来，才是真正的“性能核弹”。

---