Beware that differentiating a matrix exponential (with respect to parameters of the matrix) is not as simple as many people expect (but ChainRules.jl and hence Zygote.jl can do it). See also Differentiating random walk probability w.t.r. rate of jump - #14 by stevengj

1 Like