Closed ChrisRackauckas closed 4 years ago
I think it's something like:
if S.noiseterm
vecjacobian!(dλ, λ, p, t, S, dy=dy)
for (i, λi) in enumerate(λ)
_, back = Zygote.pullback(y, prob.p) do u, p
f(u, p, t)[:,i]
end
_,tmp2 = back(FillArrays.Fill(λi,length(λ)')
dgrad[:,i] .= vec(tmp2)
end
else
vecjacobian!(dλ, λ, p, t, S, dgrad=dgrad, dy=dy)
end
In support of this: Each of the columns of the noise matrix is scaled with dW[k]
like the drift vector is scaled with dt
. So by taking the total derivative, you expect a sum over terms which are scaled variants of what you get for an ODE adjoint.
The current implementation of the adjoints on SDEs requires diagonal noise. We should make sure to work out the more general case and add handling for that.