Data

Tools

Indexes

Applications

Experiments

Open Source Science

https://github.com/google-deepmind/pgmax

Loopy belief propagation for factor graphs on discrete variables in JAX

https://github.com/google-deepmind/pgmax

Science Score: 36.0%

This score indicates how likely this project is to be science-related based on various indicators:

  • ○
    CITATION.cff file
  • ✓
    codemeta.json file
    Found codemeta.json file
  • ○
    .zenodo.json file
  • ✓
    DOI references
    Found 1 DOI reference(s) in README
  • ✓
    Academic publication links
    Links to: arxiv.org, science.org
  • ○
    Committers with academic emails
  • ○
    Institutional organization owner
  • ○
    JOSS paper metadata
  • ○
    Scientific vocabulary similarity
    Low similarity (14.9%) to scientific vocabulary

Keywords

bayesian-inference belief-propagation jax python
Last synced: 4 months ago · JSON representation

Repository

Loopy belief propagation for factor graphs on discrete variables in JAX

Basic Info
  • Host: GitHub
  • Owner: google-deepmind
  • License: apache-2.0
  • Language: Jupyter Notebook
  • Default Branch: main
  • Homepage:
  • Size: 14.6 MB
Statistics
  • Stars: 152
  • Watchers: 8
  • Forks: 12
  • Open Issues: 1
  • Releases: 4
Topics
bayesian-inference belief-propagation jax python
Created over 3 years ago · Last pushed over 1 year ago
Metadata Files
Readme Contributing License

README.md

continuous-integration PyPI version Documentation Status

PGMax

PGMax implements general factor graphs for discrete probabilistic graphical models (PGMs), and hardware-accelerated differentiable loopy belief propagation (LBP) in JAX.

  • General factor graphs: PGMax supports easy specification of general factor graphs with potentially complicated topology, factor definitions, and discrete variables with a varying number of states.
  • LBP in JAX: PGMax generates pure JAX functions implementing LBP for a given factor graph. The generated pure JAX functions run on modern accelerators (GPU/TPU), work with JAX transformations (e.g. vmap for processing batches of models/samples, grad for differentiating through the LBP iterative process), and can be easily used as part of a larger end-to-end differentiable system.

See our companion paper for more details.

Installation | Getting started

Installation

Install from PyPI

pip install pgmax

Install latest version from GitHub

pip install git+https://github.com/deepmind/PGMax.git

Developer

While you can install PGMax in your standard python environment, we strongly recommend using a Python virtual environment to manage your dependencies. This should help to avoid version conflicts and just generally make the installation process easier.

git clone https://github.com/deepmind/PGMax.git cd PGMax python3 -m venv pgmax_env source pgmax_env/bin/activate pip install --upgrade pip setuptools pip install -e .

Install on GPU

By default the above commands install JAX for CPU. If you have access to a GPU, follow the official instructions here to install JAX for GPU.

Getting Started

Here are a few self-contained Colab notebooks to help you get started on using PGMax. We recommend running them on GPU instances:

Citing PGMax

PGMax is part of the DeepMind JAX ecosystem. If you use PGMax in your work, please consider citing our companion paper @article{zhou2022pgmax, author = {Zhou, Guangyao and Dedieu, Antoine and Kumar, Nishanth and L{\'a}zaro-Gredilla, Miguel and Kushagra, Shrinu and George, Dileep}, title = {{PGMax: Factor Graphs for Discrete Probabilistic Graphical Models and Loopy Belief Propagation in JAX}}, journal = {arXiv preprint arXiv:2202.04110}, year={2022} } and using the DeepMind JAX Ecosystem citation bibtex @software{deepmind2020jax, title = {The {D}eep{M}ind {JAX} {E}cosystem}, 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}, url = {http://github.com/google-deepmind}, year = {2020}, }

Note

This is not an officially supported Google product.

Owner

  • Name: Google DeepMind
  • Login: google-deepmind
  • Kind: organization

GitHub Events

Total
  • Issues event: 1
  • Watch event: 29
  • Push event: 3
  • Fork event: 1
Last Year
  • Issues event: 1
  • Watch event: 29
  • Push event: 3
  • Fork event: 1

Committers

Last synced: 8 months ago

All Time
  • Total Commits: 47
  • Total Committers: 4
  • Avg Commits per committer: 11.75
  • Development Distribution Score (DDS): 0.319
Past Year
  • Commits: 5
  • Committers: 2
  • Avg Commits per committer: 2.5
  • Development Distribution Score (DDS): 0.4
Top Committers
Name Email Commits
PGMax team p****v@g****m 32
Antoine Dedieu t****u@g****m 9
StannisZhou t****o@g****m 4
Stannis Zhou s****s@g****m 2
Committer Domains (Top 20 + Academic)
google.com: 2

Issues and Pull Requests

Last synced: 8 months ago

All Time
  • Total issues: 8
  • Total pull requests: 0
  • Average time to close issues: 3 months
  • Average time to close pull requests: N/A
  • Total issue authors: 8
  • Total pull request authors: 0
  • Average comments per issue: 2.25
  • Average comments per pull request: 0
  • Merged pull requests: 0
  • Bot issues: 0
  • Bot pull requests: 0
Past Year
  • Issues: 2
  • Pull requests: 0
  • Average time to close issues: about 2 hours
  • Average time to close pull requests: N/A
  • Issue authors: 2
  • Pull request authors: 0
  • Average comments per issue: 4.5
  • Average comments per pull request: 0
  • Merged pull requests: 0
  • Bot issues: 0
  • Bot pull requests: 0
View more stats
Top Authors
Issue Authors
  • hnisonoff (1)
  • gaborcs (1)
  • amrzv (1)
  • murphyk (1)
  • cuent (1)
  • hayoc (1)
  • NishanthJKumar (1)
  • ishaanmht (1)
  • lkskstlr (1)
Pull Request Authors
Top Labels
Issue Labels
Pull Request Labels

Dependencies

.github/workflows/ci.yaml actions
  • actions/checkout v2 composite
  • actions/setup-python v2 composite
  • actions/setup-python v1 composite
  • styfle/cancel-workflow-action 0.8.0 composite
.github/workflows/publish.yaml actions
  • actions/checkout v2 composite
  • actions/setup-python v1 composite
docs/requirements.txt pypi
  • IPython ==8.10
  • ipykernel ==5.3.4
  • jax >=0.2.25
  • jaxlib >=0.1.74
  • jinja2 <3.1.0
  • joblib >=1.1.0
  • jupyter >=1.0.0
  • jupytext >=1.11.3
  • matplotlib >=3.2.0
  • numba >=0.55.0
  • numpy >=1.21.0
  • pandoc ==1.0.2
  • pre-commit >=2.13.0
  • pytest >=7.0.1
  • pytest-cov >=2.12.1
  • pytest-env >=0.6.2
  • scikit-learn >=1.0.1
  • scipy >=1.2.3
  • sphinx ==3.3.0
  • sphinx-autodoc-typehints ==1.11.1
  • sphinx_rtd_theme ==0.5.0
  • sphinxcontrib-bibtex ==1.0.0
  • sphinxcontrib-katex ==0.7.1
  • tensorflow >=2.9.1
  • tensorflow_datasets >=4.6.0
  • tqdm >=4.61.0
requirements.txt pypi
  • cvxpy >=1.2.0
  • jax >=0.2.25
  • jaxlib >=0.1.74
  • joblib >=1.1.0
  • jupyter >=1.0.0
  • jupytext >=1.11.3
  • matplotlib >=3.2.0
  • numba >=0.55.0
  • numpy >=1.21.0
  • pre-commit >=2.13.0
  • pytest >=7.0.1
  • pytest-cov >=2.12.1
  • pytest-env >=0.6.2
  • scikit-learn >=1.0.1
  • scipy >=1.2.3
  • sphinx >=4.4.0
  • sphinx-book-theme >=0.3.3
  • sphinx-rtd-theme >=0.5.2
  • tensorflow >=2.9.1
  • tensorflow_datasets >=4.6.0
  • tqdm >=4.61.0
  • typing_extensions *
setup.py pypi