JIT changes the output of a function

[fbartolic]

What's going on here?

import jax.numpy as jnp
from jax import jit, random

key = random.PRNGKey(42)
key, subkey = random.split(key)

X = random.uniform(subkey, (2, 100))

def f1(X):
    X = X.at[0].set(jnp.roll(X[0], 5, axis=-1))
    return X

def f2(X):
    X = X.at[0].set(X[0]**2)
    return X

def f3(X):
    x = jnp.roll(X[0], 5, axis=-1) 
    return jnp.stack([x, X[1]])

print(jnp.all(f1(X) == jit(f1)(X))) # evaluates to False
print(jnp.all(f2(X) == jit(f2)(X))) # evaluates to True
print(jnp.all(f3(X) == jit(f3)(X))) # evaluates to True

[jakevdp]

Thanks for the report! This is quite reminiscent of #7461, and similar to that issue only appears to happen on CPU. I suspect the same mechanism is somehow at play

[jakevdp]

I tried with the most recent jax & jaxlib on a CPU backend - it appears the issue has been fixed.