buggy interaction: remat, automatic partitioning, and unsafe rbg-based RNGs

[froystig]

Description

Consider the interaction of the following features:

1. jax.remat
2. rbg or unsafe_rbg RNG implementations
3. the XLA flag xla_tpu_spmd_rng_bit_generator_unsafe
4. automatic partitioning

remat says: jax, you are free not to save the output of a random draw, instead only save the key and recompute random bits from it. The system is happy to do this because RNGs are by assumption functional and deterministic.

rbg and unsafe_rbg RNGs say: produce random bits from XLA's RngBitGenerator HLO.

The XLA flag xla_tpu_spmd_rng_bit_generator_unsafe says: random bits produced by a given key are allowed to vary based on the sharding of generated random bit arrays.

Automatic partitioning says: compiler, you are free to shard random bit arrays however makes sense, potentially sharding them differently when rematerializing them in the generated VJP.

Putting these together, random bits rematerialized (from the same random key) might be different between forward and backward passes in a VJP.

System info (python version, jaxlib version, accelerator, etc.)

0.4.27.dev with 0.4.26 jaxlib.