PGMax is a library for working with Factor Graphs in JAX. 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.
- Install the library using pip via:
pip install pgmax
- By default this installs JAX for CPU. If you'd like to use JAX with a GPU and specific CUDA version (highly recommended), follow the official instructions here.
- Clone this project's GitHub Repository
- Install Poetry by following these instructions. Note: you may need to logout and log back in after running the install command for the
poetry --version
command to work in your shell environment. - Navigate to this project's directory and activate a poetry shell via the command
poetry shell
. This creates and activates a virtual environment for you to use with this project. - Install the project's dependencies into your virtual environment with the command
poetry install
. Your environment will now contain both developer and user dependencies!- By default this installs JAX for CPU. If you'd like to use JAX with a GPU and specific CUDA version (highly recommended), follow the official instructions here.
- Do
pre-commit install
to initialize pre-commit hooks