share rng and predicate between all foiled args (#2375)
share rng and predicate between all foiled args (#2375)

This doesn't appear to be a performance disadvantage, and it makes
reading the HLO infinitely easier, as there's now only one pair of
RNG/constants per remat call.
google/jax,tensorflow/probability,tensorflow/probability,google/jax,google/jax,google/jax