jax.jit无法提升大数组原地更新性能？问题演示与求助

核心原因 #

你的测试结果符合JAX的设计特性，根本问题在于JAX和Numpy的底层机制差异：

• JAX的“原地更新”是语法糖，本质是创建新数组：
  JAX数组是不可变的，at.set并不会修改原数组的内存，而是返回一个包含修改的新数组。对于10000x10000的大型数组，每次单元素更新都要复制整个数组的内存，这是耗时的主要来源。而Numpy的b[1,1] = 1是直接修改内存中的指定位置，几乎没有额外开销。

• JIT不适合优化细粒度单元素操作：
  JAX的JIT编译器擅长优化批量、向量/矩阵级别的计算，单元素更新这种极细粒度的操作，JIT编译带来的收益完全无法抵消数组复制、设备同步的开销。加上block_until_ready()强制同步JAX的异步执行流程，进一步放大了延迟。

• 设备传输开销（若使用GPU/TPU）：
  JAX默认会将数组放到GPU/TPU上运行，每次单元素更新都涉及主机与设备之间的数据传输，这也是性能差距的重要因素；而Numpy默认在CPU内存中操作，没有这层开销。

解决方案 #

1. 批量更新替代单元素更新 #

JAX的优势在于批量操作，把所有需要更新的索引和值收集起来，一次性执行更新，避免多次数组复制。示例代码：

import jax.numpy as np
from jax import jit

node_count = 10000
a = np.zeros([node_count, node_count])

# 定义批量更新函数
def batch_update(mat, indices, vals):
    return mat.at[tuple(np.array(indices).T)].set(vals)

# JIT编译批量更新函数
batch_update_jit = jit(batch_update)

# 准备1000个更新任务
indices = [[i, i] for i in range(1000)]
vals = np.ones(1000)

# 测试批量更新性能
%timeit batch_update_jit(a, indices, vals).block_until_ready()

这种批量操作的性能会远高于多次单元素更新，能充分发挥JAX的优化能力。

2. 混合使用Numpy和JAX #

如果必须执行大量细粒度单元素修改，建议先用Numpy完成这类操作，再将数组转为JAX数组进行后续的批量计算，兼顾细粒度操作的性能和JAX的批量计算优势。

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