Odd output from jax.ops.index_update when jitted

import jax
import jax.numpy as jnp

def demo(n=8):
  fn = lambda x: jax.ops.index_update(x, slice(1, None), 1 + x[:-1])
  y = jnp.zeros(n)
  print(fn(y))
  print(jax.jit(fn)(y))
demo()

# [0. 1. 1. 1. 1. 1. 1. 1.]
# [0. 1. 2. 3. 4. 5. 6. 7.]

[hawkinsp]

We have a candidate fix for this issue that hopefully should land soon.

[hawkinsp]

Debugging progress is happening. We suspect that this isn’t a CPU-only bug, it’s just much more likely to exhibit on CPU. Watch this space!