Consider migrating to diffrax #771

Diffrax is a JAX-powered ODE solver from which the project might benefit.

Please complete the following tasks

• I have searched the documentation and help resources before reporting.
• I have searched the discussions to avoid duplicates.
• I have searched the open and closed issues to avoid duplicates.

Is your feature request related to a problem? Please describe.

The solver we currently use (SciPy) is not reliable enough, largely because our use cases are too stiff.

Describe the solution you'd like

Anything that Just Works® would be great.

Describe alternatives you've considered

Rewrite in Rust, perform some pre-equilibrations, tweak with SciPy options, etc. None of those should be as simple as the proposed one?

Additional context

The stiffest problems seem to be the ones with equilibria. Maybe we should test more stuff like that as well?