Hacker News new | ask | show | jobs
by sitkack 524 days ago
Did you use https://github.com/patrick-kidger/diffrax ?

JAX Talk: Diffrax https://www.youtube.com/watch?v=Jy5Jw8hNiAQ
2 comments
yberreby 524 days 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/ [2]: https://github.com/patrick-kidger/equinox

link
sitkack 524 days ago
Thanks, that looks neat.

Kidger's thesis is wonderful https://arxiv.org/abs/2202.02435

link
marmaduke 524 days 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

link
Iwan-Zotow 524 days ago
Not as effective as I'd like?
link
marmaduke 524 days ago
ha, yeah, thanks.
link