pgjax draws Pólya-Gamma variates as an XLA
custom call, so a Gibbs sampler that needs them stays on the accelerator. It
replaces jax.pure_callback around a host-side sampler, which pays a round trip
between device and host on every sweep and serializes under jax.pmap.
