{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/pgmax-factor-graphs-for-discrete","title":"PGMax: Factor Graphs for Discrete Probabilistic Graphical Models and Loopy Belief Propagation in JAX","arxiv_id":"2202.04110","date":"2022-02-08","proceeding":null,"authors":["Guangyao Zhou","Antoine Dedieu","Nishanth Kumar","Wolfgang Lehrach","Miguel Lázaro-Gredilla","Shrinu Kushagra","Dileep George"],"abstract":"PGMax is an open-source Python package for (a) easily specifying discrete Probabilistic Graphical Models (PGMs) as factor graphs; and (b) automatically running efficient and scalable loopy belief propagation (LBP) in JAX. PGMax supports general factor graphs with tractable factors, and leverages modern accelerators like GPUs for inference. Compared with existing alternatives, PGMax obtains higher-quality inference results with up to three orders-of-magnitude inference time speedups. PGMax additionally interacts seamlessly with the rapidly growing JAX ecosystem, opening up new research possibilities. Our source code, examples and documentation are available at https://github.com/deepmind/PGMax.","url_abs":"https://arxiv.org/abs/2202.04110v4","url_pdf":"https://arxiv.org/pdf/2202.04110v4.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"pgmax-factor-graphs-for-discrete","repo_url":"https://github.com/deepmind/pgmax","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok","spdx":"Apache-2.0"}},{"paper_slug":"pgmax-factor-graphs-for-discrete","repo_url":"https://github.com/vicariousinc/pgmax","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok","spdx":"MIT"}}],"tasks":[],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2202.04110","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2202.04110"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/vicariousinc/pgmax","reach":{"status":"ok","spdx":"MIT"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/deepmind/pgmax","reach":{"status":"ok","spdx":"Apache-2.0"}}],"summary":{"unverified":5},"by_repo_kind":{"official":{"samples":5,"ran":0,"repositories":2}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":0,"samples":[{"code_sha256_prefix":"fcab3b93541f7288","entry":"cached_property","repo":"vicariousinc/pgmax","repo_kind":"official","path":"pgmax/utils.py","file_url":"https://github.com/vicariousinc/pgmax/blob/HEAD/pgmax/utils.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"fcab3b93541f7288"}},{"code_sha256_prefix":"6823f2ba47f6e030","entry":"concatenate_var_states_for_edges","repo":"deepmind/pgmax","repo_kind":"official","path":"pgmax/factor/factor.py","file_url":"https://github.com/deepmind/pgmax/blob/HEAD/pgmax/factor/factor.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"6823f2ba47f6e030"}},{"code_sha256_prefix":"32a9aa4aab3e198d","entry":"normalize_and_clip_msgs","repo":"vicariousinc/pgmax","repo_kind":"official","path":"pgmax/infer/bp.py","file_url":"https://github.com/vicariousinc/pgmax/blob/HEAD/pgmax/infer/bp.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"32a9aa4aab3e198d"}},{"code_sha256_prefix":"8ca28643e9bba3ea","entry":"pass_var_to_fac_messages","repo":"deepmind/pgmax","repo_kind":"official","path":"pgmax/infer/bp.py","file_url":"https://github.com/deepmind/pgmax/blob/HEAD/pgmax/infer/bp.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"8ca28643e9bba3ea"}},{"code_sha256_prefix":"4d36d66fecb922f1","entry":"pass_var_to_fac_messages","repo":"vicariousinc/pgmax","repo_kind":"official","path":"pgmax/infer/bp.py","file_url":"https://github.com/vicariousinc/pgmax/blob/HEAD/pgmax/infer/bp.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"4d36d66fecb922f1"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}