JAX中列表追加式微分方程求解循环的向量化替代方案

问题背景 #

现有如下基于Python循环的微分方程求解代码（Euler法），因每次迭代的状态x依赖前一步结果，jax.vmap无法适配，希望找到JAX原生的高效实现方式：

import jax.numpy as jnp
from jax import jit

# 待积分的函数
@jit
def dxdt(t, x):
   return -x**2

# 简化的Euler积分器
@jit
def integrator(f, t, x, dt):
    return x + f(t, x) * dt

t_arr = jnp.linspace(0, 10, 100)
dt = t_arr[1] - t_arr[0]

x_list = []

# 初始化x
x = 0.

for t in t_arr:
    x_list.append(x)
    x = integrator(dxdt, t, x, dt)  # 注：原代码中参数`f`应为`dxdt`，此处修正

x_arr = jnp.array(x_list)

解决方案：使用jax.lax.scan #

对于这种状态依赖的顺序迭代，JAX提供的jax.lax.scan是最优适配方案——它能将Python循环转换为JAX可编译的高效操作，既保留迭代的状态依赖逻辑，又避免Python循环的性能开销。

改写后的代码 #

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

@jit
def dxdt(t, x):
   return -x**2

@jit
def step(carry, t):
    x_prev = carry
    # 执行Euler步更新状态
    x_next = x_prev + dxdt(t, x_prev) * dt
    return x_next, x_prev  # 返回新状态（用于下一轮迭代）和当前记录值

t_arr = jnp.linspace(0, 10, 100)
dt = t_arr[1] - t_arr[0]

# 初始状态
x_initial = 0.

# 用scan执行迭代
final_x, x_arr = lax.scan(step, x_initial, t_arr)

关键逻辑说明 #

• step函数是迭代的核心：接收上一轮的状态carry（即前一步的x）和当前时间点t，返回更新后的状态和需要记录的当前x值。
• lax.scan的参数：第一个是step函数，第二个是初始状态，第三个是迭代输入序列（即t_arr）。
• 输出结果：final_x是迭代结束后的最终状态，x_arr是所有迭代步骤的x值组成的数组，和原代码的x_arr完全一致。

为什么vmap不适用？ #

vmap仅用于处理无依赖的并行操作，每个输入的计算相互独立；而此处迭代的每一步都依赖前一步的x结果，属于强顺序依赖逻辑，vmap无法处理这类场景。

简化优化：整合积分逻辑 #

可以将积分器逻辑直接嵌入step函数，让代码更紧凑：

@jit
def step(carry, t):
    x_prev = carry
    dx = dxdt(t, x_prev)
    return x_prev + dx * dt, x_prev

内容的提问来源于stack exchange，提问作者cluelessness