This year's Deep Learning Indaba had a tutorial on diffusion models in Jax: https://github.com/deep-learning-indaba/indaba-pracs-2022/tr...