Flakiness in JAX deterministic runs

[cool-RR]

JAX's deterministic runs feature is useful to me. I write "golden tests" in which I perform a certain sequence of operations using my package. The test asserts that the results are exactly as they were before. I have a script that generates data for these golden tests whenever I make a meaningful change to my logic that causes the results to differ. I love this setup, because I get high confidence that I don't change the results of my experiments when I write code in my project.

One thing I've been noticing recently is that sometimes when I run my golden tests after a change, they fail, i.e. the results are different from the recorded results. Then when I run them again, they pass, and I can run them over and over again and they'll pass. What I think might be happening is that when I run the golden tests for the first time after a change, JAX doesn't find the compilation of my function(s) in the cache, and compiles it, while it subsequent runs it finds the compilation in the cache. I'm guessing that this causes the results to differ.

This is a nuisance to me because I end up having to run the tests at least twice every time, to check whether each failure was real or fake. This makes me lose faith in my tests.

Did anyone ever experience this before? While I could think of a few ways to solve this problem, I'll appreciate any input you may have on why this problem is happening and how to best solve it.

[pearu]

There is #28064 that discusses result differences in jitted and non-jitted programs.

Can you provide a simple reproducer of your problem?

[cool-RR]

Thanks Pearu. I always run my program with JIT so I doubt that's the explanation. Unfortunately I can't provide a simple reproducer.