Builtin differential testing mechanism in JAX

[PragmaTwice]

Hi, everyone. I recently discovered that if we develop a new PJRT plugin for JAX (different from the current XLA:CPU, XLA:CUDA, etc.), we will face some testing troubles, especially when the JAX program is extremely complex.

I thought that if there is a built-in differential testing mechanism in JAX: for example, when we want to test the correctness of a new GPU backend, we can do it by comparing the test results of the CPU and GPU (if ignore the precision difference). That is, when some JAX operators are called (such as jax.numpy.add, here we assume that jax.jit is not enabled), JAX (in pjit or pxla) can not only delegate the calculation to the new GPU backend for execution, but also delegate to CPU (or XLA:CUDA), and then compare the results at the end. Thus we can easily determine which intermediate result is not correct in the calculation (i.e. the PJRT plugin implementation has a problem), which should be meaningful when debugging the new PJRT plugin on large JAX programs.

I made some modifications to pjit.py/pxla.py in JAX to show the feasibility of these ideas, but I ran into some problems (got some device assignment errors). I was wondering if the JAX team has any thoughts or suggestions on this?

[yashk2810]

Why not just run the same function under different backends one after another? I think that approach would be simpler than having multi-dispatch built-into jit.

[PragmaTwice]

Yeah you are right that we can simply run the same jax program with different backend e.g. via different JAX_PLATFORMS. However there are several advantages if we can having this functionality built-in:

• Some JAX programs are quite large and complex (such as AlphaFold), potentially involving millions of operations (i.e., PjitFunction, if jax.jit is not enabled). When debugging issues in our PJRT backend, it would be extremely helpful if we could inspect the intermediate results of each executed operation and abort as soon as a computation error occurs. This would make it much easier to pinpoint the exact operation causing the issue, without requiring extensive modifications to users' JAX programs.
• It would help reduce false positives caused by randomness in JAX programs, such as variations in random seeds. While this can also be mitigated by modifying JAX programs to be pure functions without randomness, having built-in support would provide a more convenient and reliable solution.

Of course, we are also exploring other approaches to address these issues. However, I am particularly interested in whether you have encountered similar challenges while developing JAX. If you could share any insights or experiences, I would greatly appreciate it : )