how to dynamic_slice with DeviceArray

[L-Hugh]

@jax.jit
def test(a, b):
    return jax.lax.dynamic_slice(a, (0,), (b,))
a = jnp.array([1,2,3])
b = jnp.array(1)
test(a, b)

While tracing the function test at :1, this concrete value was not available in Python because it depends on the value of the arguments to test at :1 at flattened positions [1], and the computation of these values is being staged out (that is, delayed rather than executed eagerly).

Is there a way to do it liske this?

[jheek]

The issue here is that b affects the shape of the result returned by dynamic_slice. You can do this by declaring b as a static argument:

def test(a, b):
    return jax.lax.dynamic_slice(a, (0,), (b,))
a = jnp.array([1,2,3])
b = jnp.array(1)
jax.jit(test, static_argnums=(1,))(a, b)

Note however that test will be re-compiled if you pass in a different value for b

[L-Hugh]

Thank you! I think I get it.