2022/02/08 by Guangyao Zhou, Zhou, Guangyao, Nishanth Kumar +9
Biochemistry, Genetics and Molecular Biology · Computer Science · #Artificial Intelligence (cs.AI) #Bayesian Modeling and Causal Inference #Bioinformatics and Genomic Networks #FOS: Computer and information sciences #Gene expression and cancer classification #Machine Learning (cs.LG) #Machine Learning (stat.ML)
paper · pdf · doi:10.48550/arxiv.2202.04110
openalex publication_date 2022/02/08 · openalex created_date 2022/04/03 · openalex updated_date 2026/07/28
PGMax is an open-source Python package for (a) easily specifying discrete Probabilistic Graphical Models (PGMs) as factor graphs; and (b) automatically running efficient and scalable loopy belief propagation (LBP) in JAX. PGMax supports general factor graphs with tractable factors, and leverages modern accelerators like GPUs for inference. Compared with existing alternatives, PGMax obtains higher-quality inference results with up to three orders-of-magnitude inference time speedups. PGMax additionally interacts seamlessly with the rapidly growing JAX ecosystem, opening up new research possibilities. Our source code, examples and documentation are available at https://github.com/deepmind/PGMax.