求导与控制流

jax.grad 通过追踪您的 Python 函数来构建一个计算图，然后对该图进行求导。当您的函数包含 Python 原生的控制流语句（如 if、for 和 while）时，这种追踪机制会带来重要的影响。

核心原则如下：JAX 针对给定的一组输入形状和类型，只追踪函数一次。在此初始追踪过程中，通过控制流语句所采取的特定路径，是唯一被编译（如果使用 jit）并求导的路径。

条件语句 (if/else)

考虑一个带有简单 if 语句的函数：

import jax
import jax.numpy as jnp

def conditional_func(x):
  if x > 0:
    return x * x
  else:
    return -x

让我们尝试获取它的梯度：

grad_conditional_func = jax.grad(conditional_func)

# 当 x > 0 时会怎样？
print(grad_conditional_func(2.0))
# 输出：4.0 (x*x 的正确导数是 2x，所以 2*2=4)

# 当 x <= 0 时会怎样？
try:
    print(grad_conditional_func(-2.0))
except Exception as e:
    print(f"Error: {e}")
# 输出：错误：遇到抽象追踪器值...
# ...抽象追踪器值的真值不明确。

为什么会出现这个错误？当 JAX 使用 2.0 这样的正值追踪 conditional_func 时，条件 x > 0 评估为 True。JAX 追踪 return x * x 分支。生成的追踪计算只包含平方操作。当您随后使用 -2.0 调用 梯度 函数时，它尝试使用新输入执行这个追踪到的图（该图只知道平方操作），但原始条件 x > 0 现在涉及到抽象的 追踪器 值。Python 的 if 不知道如何处理这些抽象值，从而导致错误。追踪没有捕获到 else 分支。

解决方案：jax.lax.cond

当条件语句依赖于 JAX 正在追踪的值（例如您想对其求导的函数参数 (parameter)）时，您需要使用 JAX 特定的控制流原语。对于 if/else，这就是 jax.lax.cond。

jax.lax.cond 接受四个参数：

1. pred：布尔条件（可以从追踪值中导出）。
2. true_fun：当 pred 为真时要执行的函数。
3. false_fun：当 pred 为假时要执行的函数。
4. operand：传递给 true_fun 或 false_fun 的输入操作数。

让我们重写函数：

import jax
import jax.numpy as jnp
from jax import lax

def conditional_func_lax(x):
  return lax.cond(
      x > 0,                # 条件
      lambda operand: operand * operand,  # 真分支函数
      lambda operand: -operand,         # 假分支函数
      x                     # 操作数
  )

grad_conditional_func_lax = jax.grad(conditional_func_lax)

# 现在它在两种情况下都有效：
print(f"Gradient at x=2.0: {grad_conditional_func_lax(2.0)}")
# 输出：x=2.0 时的梯度：4.0
print(f"Gradient at x=-2.0: {grad_conditional_func_lax(-2.0)}")
# 输出：x=-2.0 时的梯度：-1.0 ( -x 的正确导数是 -1)

通过使用 lax.cond，JAX 能够理解计算图包含了 所有 潜在的分支。然后，它就能无论输入值如何，都能正确地对函数求导，并根据执行时实际采取的分支，恰当地应用链式法则。

循环 (for/while)

当循环条件或迭代次数依赖于追踪值时，也会出现类似的问题。

静态循环： 如果 Python for 循环迭代固定次数，且在追踪期间已知，JAX 通常会 展开 该循环。

def fixed_loop(x, n=3): # n 是固定的
  y = x
  for _ in range(n):
    y = y * 2
  return y

grad_fixed_loop = jax.grad(fixed_loop)
print(f"Gradient of fixed_loop at x=1.0: {grad_fixed_loop(1.0)}")
# 输出：fixed_loop 在 x=1.0 时的梯度：8.0
# (函数为 y = x * 2^3 = 8x，导数为 8)

这之所以有效，是因为 n=3 是静态的。JAX 有效地追踪了 y = x * 2; y = y * 2; y = y * 2。然而，对于大量迭代的循环，展开会非常低效，导致计算图变得很大。

动态循环： 如果迭代次数或 while 循环的条件依赖于追踪值，标准的 Python 循环会像 if 语句示例那样失败。

• 对于迭代次数动态（由追踪值确定）的 for 循环，使用 jax.lax.fori_loop。
• 对于条件依赖于追踪值的 while 循环，使用 jax.lax.while_loop。

这些函数要求您以特定的函数式方式组织循环逻辑，通常涉及一个在每次迭代中更新的循环 carry 状态。

使用 jax.lax.fori_loop 的示例：

让我们实现 y = x * 2**n，其中 n 现在是动态输入。fori_loop 接受 lower、upper、body_fun、init_val。

import jax
import jax.numpy as jnp
from jax import lax

def dynamic_loop_lax(x, n):
  # body_fun 接受 (迭代次数, 当前 carry)
  # 这里我们只需要 carry (y)。
  def loop_body(_, current_y):
    return current_y * 2

  # 从 0 到 n-1 循环，y 从 x 开始
  final_y = lax.fori_loop(0, n, loop_body, x)
  return final_y

grad_dynamic_loop_lax = jax.grad(dynamic_loop_lax, argnums=0) # 对 x 求导

# 计算 x=1.0, n=3 次迭代时的梯度
print(f"Gradient at x=1.0, n=3: {grad_dynamic_loop_lax(1.0, 3)}")
# 输出：x=1.0, n=3 时的梯度：8.0

jax.lax.while_loop 具有类似的结构，接受 cond_fun、body_fun 和 init_val。

对梯度的影响

使用 JAX 的控制流原语（lax.cond、lax.fori_loop、lax.while_loop）可确保 JAX 能够追踪包含条件逻辑或循环行为本身的计算表示。这使得 jax.grad 能够通过对前向传播过程中给定输入所采取的实际路径进行求导来计算正确的梯度。

如果您在控制流依赖于追踪值时不使用这些原语，将会出现以下情况：

1. 您在追踪或梯度计算期间很可能会遇到错误。
2. 即使它似乎对一个输入有效（因为追踪了特定的路径），梯度函数也可能对会采取不同路径的其他输入产生不正确的结果（例如零或 NaN），因为这些路径从未被追踪过。

请记住，jax.grad 对其 追踪到的 函数进行求导。如果在追踪过程中，您的部分代码由于基于 初始 追踪输入的 Python 级别控制流而被跳过，那么梯度将无法通过这些被跳过的部分。使用 jax.lax 控制流构造是使这种动态行为对 JAX 的追踪和求导机制明确的方式。