[trostaft]

Hey, diffrax is cool by the way. I've used it to do some automatically differentiate through some ODEs for optimal experimental design. Great stuff.