Returned values of jax.random.normal slightly differ between jax[cpu]<=0.3.1 and jax[cpu]>=0.3.2

[muupan]

I noticed that returned values of jax.random.normal differ between jax[cpu]<=0.3.1 and jax[cpu]>=0.3.2. I could not find any information on this difference in CHANGELOG https://github.com/google/jax/blob/main/CHANGELOG.md#jax-032-march-16-2022. What is the cause of this difference? Is there anything I can do to match the behavior of >=0.3.2 to that of <=0.3.1?

0.3.1

> JAX_DISABLE_JIT=1 .venv-jax-0.3.1-cpu/bin/python -c "import jax; print(jax.__version__, jax.default_backend(), jax.config.jax_disable_jit); print(jax.random.normal(jax.numpy.array([3944638598, 3114804386], dtype=jax.numpy.uint32)))"
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
0.3.1 cpu True
1.4294975

0.3.2

> JAX_DISABLE_JIT=1 .venv-jax-0.3.2-cpu/bin/python -c "import jax; print(jax.__version__, jax.default_backend(), jax.config.jax_disable_jit); print(jax.random.normal(jax.numpy.array([3944638598, 3114804386], dtype=jax.numpy.uint32)))"
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
0.3.2 cpu True
1.4294977

[mattjj]

Thanks for the question.

In general, compiler updates (in XLA, and/or in LLVM on CPU) might change floating point numerical results slightly. So:

1. such changes often result not from a change not in JAX proper (i.e. in this repo) but rather in some change lower-level software JAX depends on, and
2. in general you shouldn't expect JAX's floating point numerics to be identical across compiler versions.

In this case, if you haven't changed jaxlib between your two experiments, it may indeed have been a change in JAX.

We could bisect to try to pin down the change. If the behavior changes with different versions of jax but the same version of jaxlib, we should use a git bisect on the JAX repo. If the behavior instead changes with versions of jaxlib, I should run a Google-internal bisection to get a finer grained look at the changes.

Is there anything I can do to match the behavior of >=0.3.2 to that of <=0.3.1?

We don't know yet about the cause of this change so it's hard to guess, but in general since we can't guarantee perfect float stability across versions, your recourse is to pin the jax version.

Can you confirm whether your experiment was at identical jaxlib versions? If so, want to try a git bisect on the JAX repo?

[mattjj]

I think it's indeed a jaxlib change. At jax==0.3.2, I ran against two versions of jaxlib:

So I suspect this was an underlying compiler change, in XLA or LLVM. We'd need to run an internal bisection to pin it down...

[muupan]

Thanks for your quick answer! It is understandable that floating point numerics cannot be kept identical across versions.

Can you confirm whether your experiment was at identical jaxlib versions?

The versions of jaxlib are different. My result above are jax==0.3.1,jaxlib==0.3.0 and jax==0.3.2,jaxlib==0.3.2.

> JAX_DISABLE_JIT=1 .venv-jax-0.3.1-cpu/bin/python -c "import jax; import jaxlib; print(jax.__version__, jaxlib.__version__, jax.default_backend(), jax.config.jax_disable_jit); print(jax.random.normal(jax.numpy.array([3944638598, 3114804386], dtype=jax.numpy.uint32)))"
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
0.3.1 0.3.0 cpu True
1.4294975

> JAX_DISABLE_JIT=1 .venv-jax-0.3.2-cpu/bin/python -c "import jax; import jaxlib; print(jax.__version__, jaxlib.__version__, jax.default_backend(), jax.config.jax_disable_jit); print(jax.random.normal(jax.numpy.array([3944638598, 3114804386], dtype=jax.numpy.uint32)))"
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
0.3.2 0.3.2 cpu True
1.4294977

I tried to run with jax==0.3.2,jaxlib==0.3.0. The result is identical to jax==0.3.1,jaxlib==0.3.0, so it seems to be jaxlib, not jax, that causes the difference.

> JAX_DISABLE_JIT=1 .venv-jax-0.3.2-cpu-jaxlib-0.3.0/bin/python -c "import jax; import jaxlib; print(jax.__version__, jaxlib.__version__, jax.default_backend(), jax.config.jax_disable_jit); print(jax.random.normal(jax.numpy.array([3944638598, 3114804386], dtype=jax.numpy.uint32)))"
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
0.3.2 0.3.0 cpu True
1.4294975

[muupan]

Oh I just missed your additional comments and ran the same thing.

[mattjj]

How valuable would it be to you to track down the cause of the change in jaxlib? (Even if we find it, we may not be able to do anything about it...)

[muupan]

As long as it is just a slight difference (so far it is to me) and there is nothing a user can do to it, tracking down the root cause in jaxlib is not necessary to me.