diff --git a/README.md b/README.md index 51a202df..4039060d 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,7 @@ [](https://badge.fury.io/py/pgmax) [](https://results.pre-commit.ci/latest/github/vicariousinc/PGMax/master) [](https://codecov.io/gh/vicariousinc/PGMax) +[](https://pgmax.readthedocs.io/en/latest/?badge=latest) # PGMax diff --git a/docs/source/index.rst b/docs/source/index.rst index 1c15ea4f..4cdbf02b 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -1,6 +1,11 @@ PGMax Reference Documentation ============================== -PGMax is a library for working with Factor Graphs in `JAX <https://jax.readthedocs.io/en/latest/>`_. It currently provides an interface for specifying factor graphs of any type, as well as an efficient implementation of max-product belief propagation and inference on these graphs. +PGMax implements general factor graphs for probabilistic graphical models (PGMs) with discrete variables, and hardware-accelerated differentiable loopy belief propagation (LBP) in `JAX <https://jax.readthedocs.io/en/latest/>`_. + +- General factor graphs: PGMax goes beyond pairwise PGMs, and supports arbitrary factor graph topology, including higher-order factors. +- 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. + + .. toctree:: :maxdepth: 1 diff --git a/pyproject.toml b/pyproject.toml index 0b1982bc..1e04bc72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,6 +3,7 @@ name = "pgmax" version = "0.2.0" description = "Loopy belief propagation for factor graphs on discrete variables, in JAX!" authors = ["Stannis Zhou <stannis@vicarious.com>", "Nishanth Kumar <njk@csail.mit.edu>", "Miguel Lazaro-Gredilla <miguel@vicarious.com>", "Dileep George <dileep@vicarious.com>"] +maintainers = ["Stannis Zhou <stannis@vicarious.com>", "Nishanth Kumar <njk@csail.mit.edu>"] # Additional info license = "MIT" readme = "README.md" @@ -11,6 +12,7 @@ include = [ ] homepage = "https://github.com/vicariousinc/PGMax" repository = "https://github.com/vicariousinc/PGMax" +documentation = "https://pgmax.readthedocs.io/en/latest/" keywords = ["PGMax", "jax"] [tool.poetry.dependencies]