The recent work by Zhang et. al (@ruqizhang) titled A Langevin-like Sampler for Discrete Distributions proposes using a Langevin based algorithm for sampling from discrete distributions whose log probability is differentiable w.r.t to the variables. I've implemented the code that is inspired by the blackjax api here. I'm open for the code to be reviewed and send a PR.