Combining Optax with probabilistic ODE solvers

[pnkraemer]

Hey guys,

I'm a big fan of Optax! A few days back I published some probabilistic numerics JAX code (https://pnkraemer.github.io/probdiffeq/). I made an example notebook about combining Optax and probabilistic ODE solvers for parameter estimation in ODE models: link. The optimisers worked like a charm!

Wanted to share this in case differential equations or probabilistic numerics are of interest to anyone around here :)

Nico

[adam-hartshorne]

Have you thought about adding it to here?

https://github.com/n2cholas/awesome-jax