Diffrax v0.6.1

Features

• Compatibility with JAX 0.4.36.

• New solvers! Added stochastic Runge--Kutta methods for solving the underdamped Langevin equation. We now have:

  • diffrax.AbstractFosterLangevinSRK
  • diffrax.ALIGN
  • diffrax.QUICSORT
  • diffrax.ShOULD

  and these are used with the corresponding

  • diffrax.UnderdampedLangevinDriftTerm
  • diffrax.UnderdampedLangevinDiffusionTerm

  huge thanks to @andyElking for carefully implementing all of these, which was a huge technical task. (#453 and 2000 new lines of code!) See the Underdamped Langevin Diffusion example for more on how to use these.

Bugfixes

• If t0 == t1 and we have SaveAt(ts=...) then we now correctly output len(ts) copies of y0. (Thanks @dkweiss31! #488, #494)
• When using diffrax.VirtualBrownianTree on the GPU then floating point fluctuations would sometimes produce evaluations outside of the valid [t0, t1] region, which would raise a spurious runtime error. This is now fixed. (Thanks @mattlevine22! jax-ml/jax#24807, #524, #526)
• Complex fixes in SDEs (Thanks @Randl! #454)
• Improvements to errors, warnings, and some typo fixes (Thanks @lockwo @ddrous! #468, #478, #495, #530)

New Contributors

• @ddrous made their first contribution in #530

Full Changelog: v0.6.0...v0.6.1