Differentiable Diffusion Solver in JAX This repo provides code that solves the diffusion equation in 1D and 2D