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.
翻译:PGMax是一个开源的Python包,用于轻松指定离散概率图模型(PGM)为因子图,并在JAX中自动运行高效可伸缩的循环置信传播(LBP)。PGMax支持具有可处理因子的一般因子图,并利用现代加速器如GPU进行推断。相比现有的替代方案,PGMax获得了更高质量的推断结果,推断时间加速高达三个数量级。PGMax还可以与快速增长的JAX生态系统无缝交互,开启新的研究可能性。我们的源代码、示例和文档可在 https://github.com/deepmind/PGMax 上获取。