The old code has been removed to get rid of Zygote dependencies.
The solution method for the forward iteration is pretty straightforward, and just calls back into the f, g, and if available the observation equation. See the linear case and just replace with calls to the underlying functions
The solution method for the custom rrule is also relatively straightforward, but requires AD callbacks. i.e. declaring the rule with ChainRulesCore.rrule(config::RuleConfig{>:HasReverseMode}, etc. then using _, pb_f = rrule_via_ad(config, f, x[t], t, p) etc. to call the f and get the pullback to execute its rrule. and go backwards to the x[t] and p arguments.
It seems unlikely that this would have similar performance to the custom 1st and 2nd order implementations without serious AD optimization with accumulation, but it would be worth benchmarking even for the simple setups.
The old code has been removed to get rid of Zygote dependencies.
f, g,and if available the observation equation. See the linear case and just replace with calls to the underlying functionsChainRulesCore.rrule(config::RuleConfig{>:HasReverseMode},etc. then using_, pb_f = rrule_via_ad(config, f, x[t], t, p)etc. to call thefand get the pullback to execute its rrule. and go backwards to thex[t]andparguments.