This repo is *absolutely fantastic* if you want to get started with JAX! Will recommend this to all of my students!🚀 Big shout out to @phillip_lippe for putting this together!