What should we expect from jitted functions?

[Xin-yang-Liu]

Based on my understanding, the jit function should perform operation fusion to speed up and reduce numerical errors. However when I tested a simple function:

a=jax.random.normal(jax.random.PRNGKey(0), (256,))
b=jax.random.normal(jax.random.PRNGKey(1), (256,))

jax.jit(lambda a,b: a + b - b - a )(a,b)

I am expecting the results to be exact zeros, however, they are not. And the results of both the jitted function and non-jitted function are completely the same. So I wonder under which situations we can count on jax.jit to reduce the numerical errors.

[soraros]

The algebraic simplification a + b - b - a → 0 is not safe. Consider the case where a = nan.