search

patrick-kidger / diffrax

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/
Apache License 2.0
1.33k stars 122 forks source link

Interest in new solvers #435

Open lockwo opened 1 month ago

lockwo commented 1 month ago

I have two classes of solvers I'm implementing with diffrax which I thought could be good to actually just put in diffrax, but I wanted to check if there was interest in having them in diffrax before I go through the process to get them ready for a PR.

First are some high weak order SDE solvers (Fixed step weak order 2.0 for Ito SDE solvers, e.g. SIEA etc. from https://docs.sciml.ai/DiffEqDocs/stable/solvers/sde_solve/)

Second are some diffusion model specific solvers (e.g. https://arxiv.org/abs/2105.14080). These come in two flavors: solvers that are applicable to general SDEs (or general SDEs of a certain type of noise) but perform well/better on diffusion models and solvers that are specific to (reverse) diffusion processes.

Would any subset of these be interesting to have in diffrax?

patrick-kidger commented 1 month ago

Definitely!

We'd need to think about how to correctly express solvers that are specific to certain diffusion processes, I suspect. Ideally these would throw an error if misused. (That's not crucial though.)

lockwo commented 1 month ago

Awesome, I think I have 3 open PRs right now 😅 , so that should give some time to get things into a presentable shape (and think on solver representation for specific diffusion processes)