Taking gradients of a matrix exponential

Bizzi
Wed, 20 Dec 2023 16:28:50 +0000
New to Julia

I’m having trouble differentiating simple matrix exponentials. ForwardDiff and ReverseDiff can’t really do it, while Zygote’s result is type-unstable.

using Zygote, ForwardDiff, ReverseDiff

M = rand(2,2)
ReverseDiff.gradient(x -> sum(exp(x)),M) #MethodError
ForwardDiff.gradient(x -> sum(exp(x)),M) #MethodError

Zygote.gradient(x -> sum(exp(x)),M) #Works, but:
@code_warntype Zygote.gradient(x -> sum(exp(x)),M) #Type-Unstable

@code_warntype (click for more details)
Can anyone share any advice? If possible, I would like to avoid doing this with Zygote. I suppose using a linear ODE solver and differentiating it with SciMLSensitivity could work, but there must be a simpler solution to this, right?

11 posts - 6 participants