Tell SciML about Jacobian/VJP of ODEFunction

[hexaeder]

Hello,

I have a class of ODEFunctions for which I know how to exploit a lot of the internal structure for efficient calculations of Jacobians. I am a bit confused about how I may teach the SciML-universe about those. Generally, I think I have the following three use-cases:

1. Linear solve within stiff solvers (nonsingular mass matrix)
2. parameter estimation with SciMLSensitivity
3. initial state estimation with SciMLSensitivity

According to the ODEFunction docstring there are 4 relevant parameters

• jac (and jac_prototype) for the matrix jacobian
• paramjac (but no paramjac_prototype?) matrix jacobian with respect to parameters
• jvp and
• vjp

The SciMLSensitivity docs suggest to supply vjp and then use ZygoteVJP to make sure the solver uses those. However this part of the docs seem out of date as signature for the custom vjp does not match the one from the ODEFunction docstring, so I am not sure if this is up to date.

Also, shouldn’t there be something like paramvjp for the parameter estimation case? Will the vjp definition also be used inside the stiff solvers or does this only work for SciMLSensitivity and the linear solve will look for the jac? Should I define all of those options above?

An alternative option might be create a “lazy” SciMLOperator for my jac and paramjac without actually building the full jacobian.

Another completely different option that comes to mind is to define custom ChainRules rules for my function. But will those be actually used inside the package?

It seems like everything I need is there, but I am a bit overwhelmed by the different options. So please help me understand how those different approaches relate to each other. Thanks!

[ChrisRackauckas]

@hexaeder

The SciMLSensitivity docs suggest to supply vjp and then use ZygoteVJP to make sure the solver uses those. However this part of the docs seem out of date as signature for the custom vjp does not match the one from the ODEFunction docstring, so I am not sure if this is up to date.

Yes, that needs to be updated. For now you can always define an rrule directly on your f.

@hexaeder

Also, shouldn’t there be something like paramvjp for the parameter estimation case? Will the vjp definition also be used inside the stiff solvers or does this only work for SciMLSensitivity and the linear solve will look for the jac? Should I define all of those options above?

Not really, because if doing adjoints you only ever need the vjp and never need to construct hte Jacobian.

@hexaeder

An alternative option might be create a “lazy” SciMLOperator for my jac and paramjac without actually building the full jacobian.

Which is what we plan to do internally in the near future, making the vjp operation into an operator. But SciMLOperators neds a bit more work.

@hexaeder

Another completely different option that comes to mind is to define custom ChainRules rules for my function. But will those be actually used inside the package?

Yes.

[hexaeder]

Thanks for the reply Chris! I‘ll play around with rrule then and see if I can get it all running!

Is there some example for that somewhere? I have no prior experience in writing rrules but since the ode is inplace I am not exactly sure how this works, and ChainRules seem to warn about writing such rules.