{"id":28753246,"url":"https://github.com/google-deepmind/pgmax","last_synced_at":"2025-10-17T22:39:54.605Z","repository":{"id":66141140,"uuid":"539972071","full_name":"google-deepmind/PGMax","owner":"google-deepmind","description":"Loopy belief propagation for factor graphs on discrete variables in JAX","archived":false,"fork":false,"pushed_at":"2024-10-18T01:32:29.000Z","size":15327,"stargazers_count":152,"open_issues_count":1,"forks_count":12,"subscribers_count":8,"default_branch":"main","last_synced_at":"2025-06-17T00:48:08.160Z","etag":null,"topics":["bayesian-inference","belief-propagation","jax","python"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/google-deepmind.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2022-09-22T12:32:27.000Z","updated_at":"2025-06-12T17:04:59.000Z","dependencies_parsed_at":"2023-09-07T20:34:53.336Z","dependency_job_id":"4c14b8da-8b48-42a4-8f80-0b5c803eefd6","html_url":"https://github.com/google-deepmind/PGMax","commit_stats":null,"previous_names":["google-deepmind/pgmax","deepmind/pgmax"],"tags_count":4,"template":false,"template_full_name":null,"purl":"pkg:github/google-deepmind/PGMax","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2FPGMax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2FPGMax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2FPGMax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2FPGMax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/google-deepmind","download_url":"https://codeload.github.com/google-deepmind/PGMax/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2FPGMax/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":279162733,"owners_count":26117273,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-10-16T02:00:06.019Z","response_time":53,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["bayesian-inference","belief-propagation","jax","python"],"created_at":"2025-06-17T00:39:18.703Z","updated_at":"2025-10-17T22:39:54.583Z","avatar_url":"https://github.com/google-deepmind.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"[![continuous-integration](https://github.com/deepmind/PGMax/actions/workflows/ci.yaml/badge.svg)](https://github.com/deepmind/PGMax/actions/workflows/ci.yaml)\n[![PyPI version](https://badge.fury.io/py/pgmax.svg)](https://badge.fury.io/py/pgmax)\n[![Documentation Status](https://readthedocs.org/projects/pgmax/badge/?version=latest)](https://pgmax.readthedocs.io/en/latest/?badge=latest)\n\n# PGMax\n\nPGMax implements general [factor graphs](https://en.wikipedia.org/wiki/Factor_graph)\nfor discrete probabilistic graphical models (PGMs), and\nhardware-accelerated differentiable [loopy belief propagation (LBP)](https://en.wikipedia.org/wiki/Belief_propagation)\nin [JAX](https://jax.readthedocs.io/en/latest/).\n\n- **General factor graphs**: PGMax supports easy specification of general\nfactor graphs with potentially complicated topology, factor definitions,\nand discrete variables with a varying number of states.\n- **LBP in JAX**: PGMax generates pure JAX functions implementing LBP for a\ngiven factor graph. The generated pure JAX functions run on modern accelerators\n(GPU/TPU), work with JAX transformations\n(e.g. `vmap` for processing batches of models/samples,\n`grad` for differentiating through the LBP iterative process),\nand can be easily used as part of a larger end-to-end differentiable system.\n\nSee our [companion paper](https://arxiv.org/abs/2202.04110) for more details.\n\n[**Installation**](#installation)\n| [**Getting started**](#getting-started)\n\n## Installation\n\n### Install from PyPI\n```\npip install pgmax\n```\n\n### Install latest version from GitHub\n```\npip install git+https://github.com/deepmind/PGMax.git\n```\n\n### Developer\nWhile you can install PGMax in your standard python environment,\nwe *strongly* recommend using a\n[Python virtual environment](https://docs.python.org/3/tutorial/venv.html)\nto manage your dependencies. This should help to avoid version conflicts and\njust generally make the installation process easier.\n\n```\ngit clone https://github.com/deepmind/PGMax.git\ncd PGMax\npython3 -m venv pgmax_env\nsource pgmax_env/bin/activate\npip install --upgrade pip setuptools\npip install -e .\n```\n\n### Install on GPU\n\nBy default the above commands install JAX for CPU. If you have access to a GPU, \nfollow the official instructions [here](https://github.com/google/jax#pip-installation-gpu-cuda)\nto install JAX for GPU.\n\n## Getting Started\n\n\nHere are a few self-contained Colab notebooks to help you get started on using PGMax. We recommend running them on GPU instances:\n\n- [First tutorial for basic PGMax inference on an Ising model](https://colab.research.google.com/github/deepmind/PGMax/blob/master/examples/ising_model.ipynb)\n- [Advanced tutorial running inference on a Restricted Boltzmann Machine](https://colab.research.google.com/github/deepmind/PGMax/blob/master/examples/rbm.ipynb)\n- [Implementing max-product LBP](https://colab.research.google.com/github/deepmind/PGMax/blob/master/examples/rcn.ipynb)\nfor [Recursive Cortical Networks](https://www.science.org/doi/10.1126/science.aag2612)\n- [End-to-end differentiable LBP for gradient-based PGM training](https://colab.research.google.com/github/deepmind/PGMax/blob/master/examples/grid_mrf.ipynb)\n- [2D binary deconvolution](https://colab.research.google.com/github/deepmind/PGMax/blob/master/examples/pmp_binary_deconvolution.ipynb)\n- [Alternative inference using a Smooth Dual LP-MAP solver](https://colab.research.google.com/github/deepmind/PGMax/blob/master/examples/sdlp_examples.ipynb)\n\n## Citing PGMax\n\nPGMax is part of the DeepMind JAX ecosystem. If you use PGMax in your work, please consider citing our [companion paper](https://arxiv.org/abs/2202.04110)\n```\n@article{zhou2022pgmax,\n  author = {Zhou, Guangyao and Dedieu, Antoine and Kumar, Nishanth and L{\\'a}zaro-Gredilla, Miguel and Kushagra, Shrinu and George, Dileep},\n  title = {{PGMax: Factor Graphs for Discrete Probabilistic Graphical Models and Loopy Belief Propagation in JAX}},\n  journal = {arXiv preprint arXiv:2202.04110},\n  year={2022}\n}\n```\nand using the DeepMind JAX Ecosystem citation\n```bibtex\n@software{deepmind2020jax,\n  title = {The {D}eep{M}ind {JAX} {E}cosystem},\n  author = {DeepMind and Babuschkin, Igor and Baumli, Kate and Bell, Alison and Bhupatiraju, Surya and Bruce, Jake and Buchlovsky, Peter and Budden, David and Cai, Trevor and Clark, Aidan and Danihelka, Ivo and Dedieu, Antoine and Fantacci, Claudio and Godwin, Jonathan and Jones, Chris and Hemsley, Ross and Hennigan, Tom and Hessel, Matteo and Hou, Shaobo and Kapturowski, Steven and Keck, Thomas and Kemaev, Iurii and King, Michael and Kunesch, Markus and Martens, Lena and Merzic, Hamza and Mikulik, Vladimir and Norman, Tamara and Papamakarios, George and Quan, John and Ring, Roman and Ruiz, Francisco and Sanchez, Alvaro and Sartran, Laurent and Schneider, Rosalia and Sezener, Eren and Spencer, Stephen and Srinivasan, Srivatsan and Stanojevi\\'{c}, Milo\\v{s} and Stokowiec, Wojciech and Wang, Luyu and Zhou, Guangyao and Viola, Fabio},\n  url = {http://github.com/google-deepmind},\n  year = {2020},\n}\n```\n\n\n## Note\n\nThis is not an officially supported Google product.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle-deepmind%2Fpgmax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fgoogle-deepmind%2Fpgmax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle-deepmind%2Fpgmax/lists"}