如何用JAX实现常量内存下的函数值向量化求和？

我有一个函数f，接受整数参数并返回固定大小的数组。想对range(N)中的所有i计算f(i)的总和，但当N很大时，存储所有计算值会导致内存不足。

用简单的for循环可以实现常量内存求和，但执行速度很慢：

import jax.numpy as jnp
f = lambda i : i*jnp.identity(1000)+i  # 这个函数会快速占用大量内存
result = 0. # 通用初始化，适用于任何类数组函数
for i in range(N):
    result += f(i)

若使用JAX的vmap实现，又会因先计算所有值再求和而耗尽内存：

result = jnp.sum(vmap(f)(jnp.arange(N)),axis=0)

请问使用JAX实现该求和的正确向量化方式是什么？

可以使用JAX的jax.lax.scan实现常量内存且高效的迭代求和，它会把循环编译成高效的XLA操作，同时不会一次性存储所有中间结果。

实现代码 #

import jax
import jax.numpy as jnp

def f(i):
    return i * jnp.identity(1000) + i

N = 10000  # 示例大N值

# 定义scan的迭代函数：输入当前累加值和当前i，返回新的累加值和None（无需保存中间结果）
def step(carry, i):
    return carry + f(i), None

# 初始化累加值为与f(i)同形状的零数组
initial_carry = jnp.zeros_like(f(0))
# 执行scan：遍历jnp.arange(N)，累加所有f(i)
result, _ = jax.lax.scan(step, initial_carry, jnp.arange(N))

print(result.shape)  # 输出 (1000, 1000)，符合预期

原理说明 #

• jax.lax.scan是JAX中用于高效循环的核心算子，它会将循环逻辑编译为XLA可优化的计算图，性能远高于Python原生for循环。
• 它通过carry参数传递累加状态，每次迭代仅保留当前累加值和当前输入，不会存储所有f(i)的结果，因此内存占用为常量级，仅取决于f(i)返回数组的大小，和N无关。
• 相比vmap，scan避免了一次性生成所有f(i)的批量结果，完美解决大N下的内存溢出问题。

额外优化：分块并行（可选） #

如果想进一步利用并行计算，同时保持内存可控，可以结合vmap和scan做分块处理：将N分成若干块，每块用vmap批量计算后求和，再用scan累加各块的结果。这种方式平衡了并行性和内存占用：

def chunked_sum(f, N, chunk_size=1000):
    num_chunks = (N + chunk_size - 1) // chunk_size
    # 定义处理单块的函数
    def process_chunk(chunk_idx):
        start = chunk_idx * chunk_size
        end = min(start + chunk_size, N)
        return jnp.sum(jax.vmap(f)(jnp.arange(start, end)), axis=0)
    # 用scan累加各块结果
    initial_carry = jnp.zeros_like(f(0))
    result, _ = jax.lax.scan(lambda carry, idx: carry + process_chunk(idx), initial_carry, jnp.arange(num_chunks))
    return result

result = chunked_sum(f, N=10000, chunk_size=1000)

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