Skip to content

pgjax: JAX-native Pólya-Gamma augmentation for fast Bayesian sampling

pgjax: JAX-native Pólya-Gamma augmentation for fast Bayesian sampling
Link

pgjax is a JAX-native Pólya-Gamma sampler for Bayesian inference. It ports the excellent polyagamma package via the JAX foreign functions interface. This design supports construction of a full XLA graph and requires no CPU callbacks or breaking JIT loops.