Jax is super fun to use outside of ml!

Recently I had fun reimplementing an old (but still usable!) code for accelerator optics. It involved transfer matrices for a 6D phase space to second order. Most of the FORTRAN77 source code was just pages and pages of hand-differentiated 6x6x6 matrices (with quite non-trivial elements) and the plumbing to painstakingly propagate those jacobians around for fitting... all replaced with a single, magic, call to jax.grad(). Felt like cheating!

I'm also super interested in its application to modelling, e.g. projects like https://github.com/deepmodeling/jax-fem -- particularly for chaining different sorts of simulations and analysis together and getting gradients through the lot. Also quite magic!