Diffrax v0.2.2

Performance improvements

• Now make fewer vector field traces in several cases (#172, #174)

Fixes

• Many documentation improvements.
• Fixed several warnings about jax.{tree_map,tree_leaves,...} being moved to jax.tree_util.{tree_map,tree_leaves,...}. (Thanks @jacobusmmsmit!)
• Fixed the step size controller choking if the error is ever NaN. (#143, #152)
• Fixed some crashes due to JAX-internal changes (If you've ever seen it throw an error about not knowing how to rewrite closed_call_p, it's this one.)
• Fixed an obscure edge-case NaN on the backward pass, if you were using an implicit solver with an adaptive step size controller, got a rejected step due to the implicit solve failing to converge, and happened to also be backpropagating wrt the controller_state.

Other

• Added a new Kalman filter example (#159) (Thanks @SimiPixel!)
• Brownian motion classes accept pytrees for shape and dtype arguments (#183) (Thanks @ciupakabra!)
• The main change is an internal refactor: a lot of functionality has moved diffrax.misc -> equinox.internal.

New Contributors

• @jacobusmmsmit made their first contribution in #149
• @SimiPixel made their first contribution in #159
• @ciupakabra made their first contribution in #183

Full Changelog: v0.2.1...v0.2.2