Why JIT results in a wrong value?

[EngineerKhan]

I am trying to run a simple square function using jit, but it gives an output of 17.63997 (while the normal one gives 17.64). Why is this so?

[hawkinsp]

Can you give more details? What exactly did you run to get both outputs?

[EngineerKhan]

Here is the code:

def Sq(x):
    print("Running the Non-JIT function")
    ans = x*x
    return ans    

@jit
def JIT_Sq(x):
    print("Running the JIT function")
    ans = x*x
    return ans

x = 4.2
print(Sq(x))
print(JIT_Sq(x))  

[hawkinsp]

See: https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision

JAX enforces 32-bit calculations by default, and the different output comes from performing the calculation in 32-bit precision.

You would receive the same output from a NumPy float32 from Sq if you instead set:

x = np.float32(4.2)

[EngineerKhan]

But point is that why does it return correct 17.64 for non-JIT version?
Also, even if we do use 32 or 64 bit precision, answer should be 17.64!

The only way answer can come 17.6399... something is if we convert the 4.2 to 4.1999... something first. Can you clarify the complete mechanism behind it?

[jakevdp]

But point is that why does it return correct 17.64 for non-JIT version?

Sq() is just a Python function, and does not contain any JAX statements. So if you pass a Python scalar to the function it uses Python's standard 64-bit math. Had you passed a jax array, it would return the same result as the JIT code, because JAX computations default to 32-bit:

print(Sq(jnp.array([4.2])))
# [17.639997]

The only way answer can come 17.6399... something is if we convert the 4.2 to 4.1999... something first. Can you clarify the complete mechanism behind it?

This is exactly what is going on:

import numpy as np
np.set_printoptions(precision=12, floatmode='fixed')
print(np.array([4.2], dtype='float32'))
# [4.199999809265]

Why? The value 4.2 is not exactly representable in finite binary floating point. It's analogous to how representing 1 / 3 in decimal requires an infinite number of digits. There's a good explanation of this with further examples and references in this StackOverflow post.

[EngineerKhan]

Thanks a lot for this detailed answer. Now I am clear!