I think you’d need to jaxify the log probability (and negate it!) first, but definitely! I also shied away from pymc because the blog post (I think!) doesn’t cover
- multiple dimensions
- change of variables
both of which you might have to deal with if you use a real PPL!