# Tell SciML about Jacobian/VJP of ODEFunction

**URL:** https://discourse.julialang.org/t/tell-sciml-about-jacobian-vjp-of-odefunction/117986
**Category:** Modelling & Simulations
**Tags:** sciml
**Created:** [August 9, 2024, 8:14am UTC](https://discourse.julialang.org/t/tell-sciml-about-jacobian-vjp-of-odefunction/117986 "2024-08-09T08:14:21Z")
**Posts on this page:** 3
**Page:** 1

<div class="post-metadata">

### Author: ![hexaeder](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/hexaeder/32/24403_2.png) [@hexaeder](https://discourse.julialang.org/u/hexaeder)
#### Post date: [August 9, 2024, 8:14am UTC](https://discourse.julialang.org/t/tell-sciml-about-jacobian-vjp-of-odefunction/117986/1 "2024-08-09T08:14:21Z")

</div>

Hello,

I have a class of `ODEFunction`s 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](https://docs.sciml.ai/DiffEqDocs/stable/types/ode_types/#SciMLBase.ODEFunction) 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](https://docs.sciml.ai/SciMLSensitivity/stable/manual/differential_equation_sensitivities/#Manual-VJPs) 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](https://docs.sciml.ai/DiffEqDocs/stable/types/ode_types/#SciMLBase.ODEFunction), 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!

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [August 9, 2024, 6:54pm UTC](https://discourse.julialang.org/t/tell-sciml-about-jacobian-vjp-of-odefunction/117986/2 "2024-08-09T18:54:35Z")

</div>

> [@hexaeder](#):
>
> The [`SciMLSensitivity` docs](https://docs.sciml.ai/SciMLSensitivity/stable/manual/differential_equation_sensitivities/#Manual-VJPs) 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](https://docs.sciml.ai/DiffEqDocs/stable/types/ode_types/#SciMLBase.ODEFunction), 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.

---

<div class="post-metadata">

### Author: ![hexaeder](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/hexaeder/32/24403_2.png) [@hexaeder](https://discourse.julialang.org/u/hexaeder)
#### Post date: [August 10, 2024, 5:58am UTC](https://discourse.julialang.org/t/tell-sciml-about-jacobian-vjp-of-odefunction/117986/3 "2024-08-10T05:58:31Z")

</div>

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](https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/which_functions_need_rules.html#Functions-which-mutate-arrays).
