>I started looking for benchmarks (because it sounds like the kind of thing JAX would do well) Currently Jax has an in-progress stiff ODE solver that's about 2

>I started looking for benchmarks (because it sounds like the kind of thing JAX would do well) Currently Jax has an in-progress stiff ODE solver that's about 200x slower than SciPy https://github.com/google/jax/issues/3686#issuecomment-657090547 https://github.com/google/jax/issues/3686#issuecomment-65709... and SciPy (with JIT) is about 50x-100x slower than the pure Julia methods https://benchmarks.sciml.ai/html/MultiLanguage/wrapper_packages.html https://benchmarks.sciml.ai/html/MultiLanguage/wrapper_packa... so Jax has more than a little bit of a way to go.