JAX报错：数组切片索引需静态起始/终止/步长问题排查

问题核心在于JAX的索引规则：NumPy风格的ar[i:i+k]切片要求起始、终止索引为静态值，但你用vmap批量处理时，i是被追踪的动态值，触发静态索引检查报错。

直接用jax.lax.dynamic_slice替代原有切片语法即可解决，该函数专门支持动态起始位置的静态大小切片，适配vmap场景：

import jax
import jax.numpy as jnp

def get_slice(ar, k, i):
  # dynamic_slice参数：原数组、起始坐标（一维数组）、切片大小（一维数组）
  return jax.lax.dynamic_slice(ar, start_indices=[i], slice_sizes=[k])

vec_get_slice = jax.vmap(get_slice, in_axes=(None, None, 0))

arr = jnp.array([1, 2, 3, 4, 5])
result = vec_get_slice(arr, 2, jnp.arange(3))
print(result)
# 输出：[[1 2]
#        [2 3]
#        [3 4]]

关键说明 #

• dynamic_slice的start_indices支持动态值，可通过vmap批量传递不同起始位置
• slice_sizes必须是静态值，符合JAX对数组形状静态可知的要求（JIT编译不支持动态大小数组）
• 替换后vmap能正常追踪动态起始索引，绕过静态切片检查限制

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