Strided indexing turns into gather

[jewillco]

Description

Syntax such as a[1::100] where a is a JAX tensor inside a JIT appears to turn into a gather operation rather than a strided slice, at least on TPU. This is inefficient and jax.lax.slice already supports strides.

System info (python version, jaxlib version, accelerator, etc.)

N/A

[jakevdp]

Thanks for the report! Lowering contiguous slices to lax.slice was an optimization we made a while ago, and at the time we scoped the problem to contiguous slices for simplicity. Adding strided slices to the logic would require adding support for them in this utility:

jax/jax/_src/numpy/lax_numpy.py

Line 5490 in 1cb6971

def _attempt_rewriting_take_via_slice(arr: Array, idx: Any, mode: str | None) -> Array | None:

[jewillco]

It seems like something that should work if normal slices work. It is surprising that a syntax supported by both Python indexing/NumPy and a feature supported by JAX doesn't work efficiently.

[jakevdp]

Agreed, that's why I marked this as a performance-related enhancement.