Not at all -- I was very happy to see some Equinox code out in the wild!

I think the main slowdown was doing the model updates out of the JIT region. JAX-without-JIT is ridiculously slow.