Odd output from jax.ops.index_update when jitted

[yovadia]

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.]

[zhangqiaorjc]

I verified that with

with api.disable_jit():
  print(jax.jit(fn)(y))

the results match.

Maybe something wrong with our jit.

[mattjj]

(Hey Yaniv!)

I suspect this is an XLA:CPU bug. It doesn't reproduce on TPU:

(It seems like a plausible bug too: inputs and outputs are incorrectly aliased so we're e.g. writing the new x[1] before reading it for the update to x[2].)

[mattjj]

It also doesn't repro on GPU.

[zhangqiaorjc]

Filed bug with XLA:CPU

[Edenhofer]

I think the issue I just wanted to file is very much related to this one. In a large model with wrong JVP and VJP I was able to trim down the code to wrong index_update calls:

from jax import numpy as jnp
from jax import jit


def _debug_twolog_int(x):
    from jax.ops import index, index_update

    twolog = jnp.zeros((2 + x.shape[1], ))

    # Required as otherwise does not break
    twolog = index_update(twolog, ((0, 1), ), 0.)
    # assert jnp.all(twolog == 0.)  # Always true if not compiled
    # Does not work! wtf?!!!!
    # JAX does not seem to like `twolog[1:-1]`, though slices w/o negative
    # indices seem to work
    twolog = index_update(twolog, index[2:], twolog[0:-2] + x[0])
    # twolog = index_update(twolog, index[2:], x[0])  # Works
    return twolog


r = jnp.repeat(jnp.arange(12).reshape((1, -1)), 2, axis=0)
print("W/O JIT:", _debug_twolog_int(r))
print("W   JIT:", jit(_debug_twolog_int)(r))
print("DIFF 01:", jit(_debug_twolog_int)(r) - _debug_twolog_int(r))

Changing any line in the above resolves the issue. However, it is pretty easy to break it again, e.g. via

def _debug_yet_another_twolog_int(x):
    from jax.ops import index, index_update

    twolog = jnp.zeros((2 + x.shape[1], ))

    twolog = index_update(twolog, index[0], 0.)
    twolog = index_update(twolog, index[1], 0.)
    # assert jnp.all(twolog == 0.)  # Always true if not compiled
    # twolog = index_update(twolog, index[2:], x[0])  # Works
    twolog = index_update(twolog, index[2:], x[1])
    # Does not work! wtf?!!!!
    # JAX does not seem to like `twolog[1:-1]`, though slices w/o negative
    # indices seem to work
    twolog = index_update(twolog, index[2:], twolog[1:-1] + x[0])
    return twolog


r = jnp.repeat(jnp.arange(12).reshape((1, -1)), 2, axis=0)
print(
    "DIFF 02:",
    jit(_debug_yet_another_twolog_int)(r) - _debug_yet_another_twolog_int(r)
)

[Edenhofer]

Prior versions of JAX and jaxlib such as "jax[cpu]==0.2.13" and "jaxlib==0.1.65" seem to be affected as well. Note, I did not perform a proper git bisect but instead just looped over pip install calls and executed the above lines in a script.

[yovadia]

Thanks for looking into this! (hi Matt!)

[mattjj]

For any googlers reading this, the internal P1 bug is b/195462810. We'll update this thread when progress is made on the bug.

The best workaround I can think of for now is to use GPU or TPU backends.

[Edenhofer]

I am terribly sorry for bumping this thread even though you said you would update it upon any progress.

However, since this is blocking part of the development on our side, the workaround being impossible (our problems do not fit on a GPU), and considering that seemingly no progress has been made, I would like to ask whether the priority of the XLA issue can be increased. The error seems quite severe to me and can have dangerous silent breaking effects.

[jakevdp]

Hi @Edenhofer - unfortunately there has not been any progress on this as of yet. It's turned out to be a difficult issue to pin down.

[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!

[hawkinsp]

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

[Edenhofer]

I can confirm that updating from jaxlib==0.1.71 to jaxlib==0.1.72 resolves the issue I had with index_update.

[hawkinsp]

@Edenhofer Yes, that's right, jaxlib 0.1.72 contains the fix for this issue. Sorry this took a while to track down!