Looking for feedback on a MCMC sampler from a newbie

Hello !

This is a crosspost from PyMC discourse:

I am completely new to MCMC and wrote the idea of an algorithm in Jax and it seems to perform quite well. So I am looking for feedback if you want to try it for your posterior sampling.

Thank you. :)