Unimplemented primitive in Pallas GPU lowering: dynamic_slice

[chaoming0625]

Description

I wrote an operator using the following syntax:

    def body(j):
      def true_fn():
        y_ref[ind_ref[j, i_start: i_end]] += 1.0
      jax.lax.cond(bool_vec[j], true_fn, lambda: None)
      return j + 1

    jax.lax.while_loop(lambda j: j < length, body, 0)

But I got the error:

Unimplemented primitive in Pallas GPU lowering: dynamic_slice.

I do not know why a GPU could not handle it.

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

jax:    0.4.34
jaxlib: 0.4.34
numpy:  1.26.4
python: 3.11.5 (main, Sep 11 2023, 13:54:46) [GCC 11.2.0]
jax.devices (1 total, 1 local): [CudaDevice(id=0)]
process_count: 1
platform: uname_result(system='Linux', node='niplab-ubuntu22-0', release='6.8.0-40-generic', version='#40~22.04.3-Ubuntu SMP PREEMPT_DYNAMIC Tue Jul 30 17:30:19 UTC 2', machine='x86_64')

[superbobry]

Looking at triton-lang/triton#656, we will not be able to support this in Pallas GPU, because Triton itself doesn't support dynamic slices.

[superbobry]

Actually, thinking about this a bit more, maybe we can support dynamic_slice in a restricted way. Can you share the full reproducer, please? The snippet above is missing a few definitions.

[ayaka14732]

There are issues for slice #19214 and TPU #18897 as well.