Minimal implementation of diffusion models and stochastic interpolants, along with associated code for the numerical experiments in "Single hidden layer diffusion models provably learn simple low-dimensional structure" (https://arxiv.org/pdf/2410.11275), written in jax.
Launcher scripts under slurm_scripts can be used to reproduce the exact experiments of the paper.
The code is entirely self-contained and can be used as a basis for more complex implementations of interpolants and diffusion models.