Memory corruption with JAX>=0.2.27

In the testing suite for Veros, we test for consistency with a Fortran reference code via f2py. For this, we essentially do this:

1. Do all variable setup in Python / JAX
2. Clone state to Fortran by doing getattr(fortran_module, var_name)[...] = getattr(python_module, var_name)
3. Run Python function that is being tested
4. Run corresponding Fortran function
5. Compare

Starting with JAX 0.2.27 I noticed that one test fails sometimes (~ every 3rd run) because the Fortran result has some NaNs sprinkled into it. I do not see this with older JAX versions or when using NumPy instead of JAX, so I have to believe that the JAX / XLA runtime is somehow meddling with the memory used by the Fortran routine. I have verified that what is being copied to Fortran (in step 2) is consistent between backends and does not trigger NaN outputs by itself. I have only tested this on CPU.

I’m at a loss how to debug this further. Maybe someone has a hunch what changed in 0.2.27 that could cause this?

[hawkinsp]

@soraros As I understand it, there’s an old out-of-tree flang that is more complete but depends on a really old LLVM fork, and its in-tree successor (also flang) that can’t yet generate code. The classic situation: choose either “deprecated” or “not ready yet”. But it looks like progress is happening!

[hawkinsp]

Thanks, that’s great, I can run that. I bisected it to the commit that enabled MLIR (https://github.com/google/jax/commit/5801079a4b28874f607515c1131e75caacef5e39).

You can work around the problem even with a newer jax by setting JAX_ENABLE_MLIR=0, so hopefully that unblocks you.

I’m hoping to remove the non-MLIR path sometime soon, but not until all the bugs have been shaken out.