Neural ODEs on manifolds

[salbert83]

How do you initialize a ManifoldODEProblem so calculations are not done in place? In place calculations are problematic for AD tools like Zygote. Also, what AD do you recommend for pullbacks? Enzyme? I have a problem where defining a Neural ODE on a manifold might be useful, hence these questions. Thank you.

[kellertuer]

Do you maybe have an MWE that fails with a concrete error message? I would like to help, but with the whole JuliaManifolds group only being 2-3 developers – I can not check the whole repository for potential in-place things we could potentially change nor do I have much experience with Neural ODEs (to be precise: I have zero experience with that).

[mateuszbaran]

I haven't coded a not-in-place variant of ManifoldODEProblem yet. Manifolds.jl has some support for not-in-place computations but in-place is the default because it makes more sense in most contexts. However the interface is flexible and the default can be overridden as needed (see Circle for an example).

I intended to add neural ODEs on manifolds at some point but the library is still quite far from that. I can help if you want to work in that direction. For Zygote and reverse mode AD in general, I'd expect it would make more sense to use adjoint sensitivity analysis instead of differentiating through the solver? That would also let you keep the solver mutating.