Can JIT change values by working with a different precision?

[onnoeberhard]

I am facing a situation where a function returns different values when JITed vs. pure Python (it is difficult to make a minimum working example). When examining, I noticed that the reason might be related to the float32 precision. The following code strikes me as odd:

>>> a = jnp.array(0.006963793188333511, dtype=jnp.float32)
>>> b = jnp.array(0.006963792722672224, dtype=jnp.float32)
>>> a
DeviceArray(0.00696379, dtype=float32)
>>> b
DeviceArray(0.00696379, dtype=float32)
>>> a == b
DeviceArray(False, dtype=bool)

Is this behaviour to be expected? (I would have thought that the repr DeviceArray(0.00696379, dtype=float32) cannot refer to two different values).

The main reason I am interested is that the JITed function produces a and the pure Python version of the function produces b (no difference except for removing the @jit decorator). I am running jax version 0.3.16 and jaxlib version 0.3.15 on CPU.

[jakevdp]

Two float values having the same repr does not necessarily mean they contain the same bits. You can adjust the representation using the set_printoptions function, which JAX inherits from numpy. For example:

>>> jnp.set_printoptions(precision=10)
>>> a
DeviceArray(0.006963793, dtype=float32)
>>> b
DeviceArray(0.0069637927, dtype=float32)

As for why JIT might lead to slightly different floating point outputs, this is covered in the FAQ: https://jax.readthedocs.io/en/latest/faq.html#jit-changes-the-exact-numerics-of-outputs

Hope that helps!

[onnoeberhard]

It does, thanks a lot!