The other problem you will run into is that you are mutating your matrices in-place. This is easy enough to fix. For instance, the line:
r ./= rin[reactiondecay]
r /= rin[reactiondecay]
transition function would need to be re-written in the same manner. That might not be straightforward in the case of your
T transition matrices, and it may be easier to derive and define the adjoints by hand.
One trick I’ve used is to calculate the matrix construction with
ForwardDiff, and then use the resulting Jacobian to calculate the vector’-Jacobian product, like so:
forward(params) = ... # constructs a matrix T from params
Zygote.@adjoint function forward(params)
T = forward(params)
return T, function (params_bar)
J = ForwardDiff.jacobian(forward, params)
# backward pass is vector-Jacobian product
vjp = params_bar' * J
It may not be the most efficient way to calculate things, but I trust it more than doing it by hand.