3 ms·

Did you use https://github.com/patrick-kidger/diffrax https://github.com/patrick-kidger/diffrax ? JAX Talk: Diffrax https://www.youtube.com/watch?v=Jy5Jw8hNiA

• marmaduke 2y ago
  no, wrote it by hand for use with my own Heun implementation, since it’s for use within stochastic delayed systems. jax is fun but as effective as i’d like for CPU
  • Iwan-Zotow 2y ago
    Not as effective as I'd like?
    • marmaduke 2y ago
      ha, yeah, thanks.
• yberreby 2y ago
  Anecdotally, I used diffrax (and equinox) throughout last year after jumping between a few differential equation solvers in Python, for a project based on Dynamic Field Theory [1]. I only scratched the surface, but so far, it's been a pleasure to use, and it's quite fast. It also introduced me to equinox [2], by the same author, which I'm using to get the JAX-friendly equivalent of dataclasses. `vmap`-able differential equation solving is really cool. [1]: https://dynamicfieldtheory.org/ https://dynamicfieldtheory.org/ [2]: https://github.com/patrick-kidger/equinox https://github.com/patrick-kidger/equinox
  • sitkack 2y ago
    Thanks, that looks neat. Kidger's thesis is wonderful https://arxiv.org/abs/2202.02435 https://arxiv.org/abs/2202.02435