I’m not sure if this can do what you need, but if you can reformulate your problem in JAX (e.g. through the minimize or custom_root, then you can use this Op. I formulated it for solving ODEs, but it should work with any JAX-friendly function. You still can’t use NUTS for some unkown reason (you should be able to, but there’s a bug somewhere), but you can use ADVI, and probably Metropolis.