Skip to content

Latest commit

 

History

History
20 lines (12 loc) · 1.1 KB

README.md

File metadata and controls

20 lines (12 loc) · 1.1 KB
Raw

DMALAX - Discrete Metropolis-Adjusted Langevin Algorithm in JAX

This repository contains the code for the JAX based implementation of the work by Zhang et.al 2022 titled A Langevin-like Sampler for Discrete Distributions. It's design is heavily insipred by blackjax even borrowing some code from the api. (I implemented this code in a separate repo as part of learning how samplers for discrete distributions work, and I plan to send a PR to the official blackjax repo 🤞)

Usage

Please check the notebooks in examples directory for how to use the kernel

Todo

  • Extend the kernel for Categorical distributions. Currently only binary-valued distributions are supported
  • Add more example notebooks that implement:
    • Potts Model
    • Restricted-Boltzmann Machine (RBM) Model
    • Bayesian Neural Network (BNN)

Other versions

You can find the PyTorch implementation of the paper written by the authors themselves here -> https://github.com/ruqizhang/discrete-langevin/