如何在JAX中通过多函数映射加速两类for循环？

一、带状态更新的循环（mass_models 迭代） #

针对需要逐次更新xs/ys的循环，jax.lax.scan是合适的替代方案。核心思路是将模型方法、kwargs、etas作为扫描的输入序列，每次迭代用当前状态（xs/ys）和对应位置的参数完成计算，返回更新后的状态。

实现代码 #

import jax
import jax.numpy as jnp

# 提前整理序列参数
model_alphas = [m.alpha for m in mass_models]
kwargs_seq = kwargs
etas = jnp.array(etas)  # 转为JAX数组（若原数据不是）

def scan_step(carry, inputs):
    xs, ys = carry
    model_alpha, kw, eta = inputs
    # 对应原循环的单次计算逻辑
    dx, dy = model_alpha(xs[j], ys[j], kwargs=kw)
    weights = eta.reshape(-1, 1)
    new_xs = xs - weights * dx
    new_ys = ys - weights * dy
    return (new_xs, new_ys), None  # 返回新状态，无额外输出

# 初始化状态，打包输入序列
init_carry = (xs, ys)
scan_inputs = (model_alphas, kwargs_seq, etas)

# 执行扫描，获取最终状态
final_carry, _ = jax.lax.scan(scan_step, init_carry, scan_inputs)
final_xs, final_ys = final_carry

说明 #

• jax.lax.scan会将循环编译为XLA优化操作，消除Python循环的开销；
• 输入序列的长度需与迭代次数N一致，scan会自动按顺序匹配每个迭代的参数；
• 状态（xs/ys）会在迭代间传递更新，完全复现原循环的逻辑。

二、无状态批量映射（light_models 迭代） #

原列表推导式的逻辑是对每个索引独立计算，可改用jax.lax.map实现，它专门处理这种无状态的批量映射场景，性能优于Python列表推导式。

实现代码 #

# 整理参数序列
model_sbs = [m.surface_brightness for m in light_models]
xs_slices = xs  # xs为(N, ...)形状的JAX数组，map会自动按第一轴拆分
ys_slices = ys
kwargs_seq = kwargs
# 若pixels坐标元素形状一致，转为JAX数组；否则保留列表
pixels_x_seq = jnp.array(pixels_x_coord) if all(len(p) == len(pixels_x_coord[0]) for p in pixels_x_coord) else pixels_x_coord
pixels_y_seq = jnp.array(pixels_y_coord) if all(len(p) == len(pixels_y_coord[0]) for p in pixels_y_coord) else pixels_y_coord

def map_func(inputs):
    model_sb, x, y, kw, px, py = inputs
    return model_sb(x, y, kw, pixels_x_coord=px, pixels_y_coord=py)

# 打包所有参数序列，执行map计算
map_inputs = (model_sbs, xs_slices, ys_slices, kwargs_seq, pixels_x_seq, pixels_y_seq)
results = jax.lax.map(map_func, map_inputs)

说明 #

• jax.lax.map会自动遍历所有参数序列的对应元素，并行完成计算后堆叠结果，效果与jnp.stack的列表推导式完全一致；
• 若pixels_x_coord/pixels_y_coord的元素形状不一致，可直接传入Python列表，JAX会自动处理；
• 编译后的map操作避免了Python解释器的循环开销，在大N场景下性能提升显著。

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