PGMax is an open-source Python package for easy specification of discrete Probabilistic Graphical Models (PGMs) as factor graphs, and automatic derivation of efficient and scalable loopy belief propagation (LBP) implementation in JAX. It supports general factor graphs, and can effectively leverage modern accelerators like GPUs for inference. Compared with existing alternatives, PGMax obtains higher-quality inference results with orders-of-magnitude inference speedups. PGMax additionally interacts seamlessly with the rapidly growing JAX ecosystem, opening up exciting new possibilities. Our source code, examples and documentation are available at https://github.com/vicariousinc/PGMax.

翻译：PGMax是一个开放源码的Python包件,方便地将离散概率图形模型(PGMS)作为要素图解,并在JAX中自动生成高效和可扩展的循环信仰传播(LBP)实施。它支持一般要素图,并能够有效地利用像GPUs这样的现代加速器进行推理。与现有的替代方法相比,PGMax获得了质量更高的推论结果,以命令的重量推断速度加速。PGMax与快速增长的JAX生态系统无缝地互动,开辟了令人振奋的新可能性。我们的源代码、示例和文件可以在https://github.com/vicentinc/PGMax上查阅。