JAX
discrete-diffusion
flax