为何JAX的split()首次调用速度极慢？如何优化？

首次调用jnp.split慢的原因 #

你的推测完全正确：首次调用耗时极长是JAX的JIT（即时编译）编译开销导致的，后续调用变快则是因为JAX缓存了编译后的XLA代码。

JAX的绝大多数数组操作默认会被JIT编译，首次执行时需要将Python代码转换为优化的XLA（加速线性代数）机器码，这个过程包含形状推导、代码生成、硬件适配优化等步骤。对于jnp.split这类返回多个数组的操作，编译开销会更显著——因为要为每个分割片段生成对应的处理逻辑。当后续调用使用相同形状的输入数组和分割索引时，JAX会直接复用之前缓存的编译结果，因此执行速度大幅提升。

避免性能损耗的优化方案 #

1. 显式封装JIT编译函数 #

将分割操作封装为函数并用jax.jit装饰，显式触发编译并缓存结果，后续调用直接复用：

import jax
import jax.numpy as jnp
from timeit import default_timer as timer

@jax.jit
def split_array(array, split_indices):
    return jnp.split(array, split_indices)

# 初始化数据
array = jnp.ones(5000)
segment_size = 10
split_indices = jnp.arange(segment_size, array.shape[0], segment_size)

# 首次调用触发编译（提前消耗编译开销）
start = timer()
segments = split_array(array, split_indices)
end = timer()
print(f'首次调用: {end - start:0.2f} s')

# 后续调用复用缓存，速度极快
for k in range(4):
    start = timer()
    segments = split_array(array, split_indices)
    end = timer()
    print(f'调用 {k+1}: {end - start:0.2f} s')

2. 用形状变换替代jnp.split（等长分割场景） #

如果是固定段长的分割（最后一段可短于段长），可以用reshape结合切片的方式替代jnp.split，这种操作的编译开销更低：

import jax.numpy as jnp

array = jnp.ones(5000)
segment_size = 10

# 场景1：数组长度能被段长整除
segments = array.reshape(-1, segment_size)  # 直接得到二维数组，每行对应一个片段

# 场景2：数组长度不能被段长整除（比如5003个元素）
array = jnp.ones(5003)
full_segment_count = array.shape[0] // segment_size
full_segments = array[:full_segment_count * segment_size].reshape(-1, segment_size)
last_segment = array[full_segment_count * segment_size:]

# 拼接为片段列表（按需转换）
segments_list = [full_segments[i] for i in range(full_segment_count)] + [last_segment]

3. 提前预热编译 #

如果程序初始化阶段有空闲，可以提前调用一次分割操作，把编译开销前置，避免业务逻辑执行时出现延迟：

import jax.numpy as jnp

array = jnp.ones(5000)
segment_size = 10
split_indices = jnp.arange(segment_size, array.shape[0], segment_size)

# 预热：首次调用消耗编译开销
_ = jnp.split(array, split_indices)

# 后续业务逻辑中的调用直接复用缓存
# ...（你的业务代码）

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