How do I pass a subset of a tracer to a function?

[andreok]

I'm trying to pass a subset of a tracer to a function using .at[].
My function requires just a subset of the values:

import jax
import jax.numpy as jnp

@jax.jit
def func(x):
  return x

len = 3

case1 = jnp.arange(0, len) == 1

x = jnp.zeros(len, dtype=jnp.float64)
func(x.at[case1])

The code works just fine without the jit decorator.
However, it gives me the following error when jit is enabled:

TypeError: Cannot interpret value of type <class 'jax._src.numpy.array_methods._IndexUpdateRef'> as an abstract array; it does not have a dtype attribute

How should I do that?

[jakevdp]

x.at[case1] does not return a subset of an array, it returns a helper that lets you generate a scatter operation (see the jnp.ndarray.at docs).

Perhaps what you want is func(x[case1])?