# Taking gradients of a matrix exponential

**URL:** <https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865>\
**Category:** New to Julia\
**Tags:** differentiation, gradient\
**Created:** [December 20, 2023, 4:28pm UTC](https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865 "2023-12-20T16:28:50Z")\
**Posts on this page:** 11\
**Page:** 1

<div class="post-metadata">

**Author:** ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)\
**Post date:** [December 20, 2023, 4:28pm UTC](https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865/1 "2023-12-20T16:28:50Z")

</div>

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

```julia
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**
>
> ```julia
> MethodInstance for Zygote.gradient(::var"#13#14", ::Matrix{Float64})
> from gradient(f, args...) @ Zygote C:\Users\55619\.julia\packages\Zygote\YYT6v\src\compiler\interface.jl:95
> Arguments
> #self#::Core.Const(Zygote.gradient)
> f::Core.Const(var"#13#14"())
> args::Tuple{Matrix{Float64}}
> Locals
> @_4::Int64
> grad::Union{Nothing, Tuple}
> back::Zygote.var"#75#76"{Zygote.Pullback{Tuple{var"#13#14", Matrix{Float64}}, var"#s178"}} where var"#s178"<:Tuple{Union{Zygote.ZBack{ChainRules.var"#exp_pullback#2042"{Tuple{Int64, Int64, Vector{Float64}, Vector{Float64}, Int64, Vector{Matrix{Float64}}, Matrix{Float64}, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Vector{Matrix{Float64}}}, Matrix{Float64}, Matrix{Float64}}}, Zygote.ZBack{ChainRules.var"#exp_pullback_hermitian#2041"{Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}, Vector{Float64}}, LinearAlgebra.Symmetric{Float64, Matrix{Float64}}, LinearAlgebra.Hermitian{Float64, Matrix{Float64}}}}}, Zygote.var"#2991#back#766"{Zygote.var"#760#764"{Matrix{Float64}}}}
> y::Float64
> Body::Union{Nothing, Tuple}
> 1 ─ %1 = Zygote.pullback::Core.Const(ZygoteRules.pullback)
> │ %2 = Core.tuple(f)::Core.Const((var"#13#14"(),))
> │ %3 = Core._apply_iterate(Base.iterate, %1, %2, args)::Tuple{Float64, Zygote.var"#75#76"{Zygote.Pullback{Tuple{var"#13#14", Matrix{Float64}}, var"#s178"}} where var"#s178"<:Tuple{Union{Zygote.ZBack{ChainRules.var"#exp_pullback#2042"{Tuple{Int64, Int64, Vector{Float64}, Vector{Float64}, Int64, Vector{Matrix{Float64}}, Matrix{Float64}, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Vector{Matrix{Float64}}}, Matrix{Float64}, Matrix{Float64}}}, Zygote.ZBack{ChainRules.var"#exp_pullback_hermitian#2041"{Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}, Vector{Float64}}, LinearAlgebra.Symmetric{Float64, Matrix{Float64}}, LinearAlgebra.Hermitian{Float64, Matrix{Float64}}}}}, Zygote.var"#2991#back#766"{Zygote.var"#760#764"{Matrix{Float64}}}}}
> │ %4 = Base.indexed_iterate(%3, 1)::Core.PartialStruct(Tuple{Float64, Int64}, Any[Float64, Core.Const(2)])
> │ (y = Core.getfield(%4, 1))
> │ (@_4 = Core.getfield(%4, 2))
> │ %7 = Base.indexed_iterate(%3, 2, @_4::Core.Const(2))::Core.PartialStruct(Tuple{Zygote.var"#75#76"{Zygote.Pullback{Tuple{var"#13#14", Matrix{Float64}}, var"#s178"}} where var"#s178"<:Tuple{Union{Zygote.ZBack{ChainRules.var"#exp_pullback#2042"{Tuple{Int64, Int64, Vector{Float64}, Vector{Float64}, Int64, Vector{Matrix{Float64}}, Matrix{Float64}, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Vector{Matrix{Float64}}}, Matrix{Float64}, Matrix{Float64}}}, Zygote.ZBack{ChainRules.var"#exp_pullback_hermitian#2041"{Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}, Vector{Float64}}, LinearAlgebra.Symmetric{Float64, Matrix{Float64}}, LinearAlgebra.Hermitian{Float64, Matrix{Float64}}}}}, Zygote.var"#2991#back#766"{Zygote.var"#760#764"{Matrix{Float64}}}}, Int64}, Any[Zygote.var"#75#76"{Zygote.Pullback{Tuple{var"#13#14", Matrix{Float64}}, var"#s178"}} where var"#s178"<:Tuple{Union{Zygote.ZBack{ChainRules.var"#exp_pullback#2042"{Tuple{Int64, Int64, Vector{Float64}, Vector{Float64}, Int64, Vector{Matrix{Float64}}, Matrix{Float64}, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Vector{Matrix{Float64}}}, Matrix{Float64}, Matrix{Float64}}}, Zygote.ZBack{ChainRules.var"#exp_pullback_hermitian#2041"{Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}, Vector{Float64}}, LinearAlgebra.Symmetric{Float64, Matrix{Float64}}, LinearAlgebra.Hermitian{Float64, Matrix{Float64}}}}}, Zygote.var"#2991#back#766"{Zygote.var"#760#764"{Matrix{Float64}}}}, Core.Const(3)])
> │ (back = Core.getfield(%7, 1))
> │ %9 = Zygote.sensitivity(y)::Core.Const(1.0)
> │ (grad = (back)(%9))
> │ %11 = Zygote.isnothing(grad)::Bool
> └── goto #3 if not %11
> 2 ─ return Zygote.nothing
> 3 ─ %14 = Zygote.map(Zygote._project, args, grad::Tuple)::Union{Tuple{}, Tuple{Any}}
> └── return %14
> 
> ```

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?

---

<div class="post-metadata">

**Author:** ![cortner](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cortner/32/204_2.png) [@cortner](https://discourse.julialang.org/u/cortner)\
**Post date:** [December 20, 2023, 8:51pm UTC](https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865/2 "2023-12-20T20:51:31Z")

</div>

Adjoint for matrix exp is fairly straightforward and indeed it is [implemented in ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl/blob/ae37562ea1f16816a0d8fff24e0aca6cd594a40f/src/rulesets/LinearAlgebra/matfun.jl#L117)

ReverseDiff and ForwardDiff however won’t use ChainRules. They both need to be told how to differentiate that function. I didn’t check this but I’m guessing there is some non-Julia in the back that prevents them from doing it out of the box. ReverseDiff.jl [can be told to use Chainrules.](https://github.com/JuliaDiff/ReverseDiff.jl/blob/b669294586b0789ab10bbda223543bcec64cf8aa/src/macros.jl#L299) Maybe that’s the best option.

I’m more concerned about the type instability when you use Zygote. I can reproduce this. Maybe somebody should comment whether this is a Zygote or a ChainRules issue. I’m guessing it is Zygote, since calling `rrule` directly is type stable:

```julia
using ChainRules, LinearAlgebra
A = randn(3,3)
expA, pb = ChainRules.rrule(exp, A)
@code_warntype pb(A)

```

P.S.: I’m assuming your example is just a MWE. Otherwise I would say just implement it yourself, possibly using `ChainRules.rrule`.

---

<div class="post-metadata">

**Author:** ![Elrod](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/elrod/32/22461_2.png) [@Elrod](https://discourse.julialang.org/u/Elrod)\
**Post date:** [December 20, 2023, 10:33pm UTC](https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865/3 "2023-12-20T22:33:11Z")

</div>

I did write this blog post comparing a few different matrix exponential implementations using `ForwardDiff.jl`, but odds are the `ChainRules` are a better approach:  
[https://spmd.org/posts/multithreadedallocations/](https://spmd.org/posts/multithreadedallocations/)

---

<div class="post-metadata">

**Author:** ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)\
**Post date:** [December 21, 2023, 2:18pm UTC](https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865/4 "2023-12-21T14:18:52Z")

</div>

> [@cortner](#):
>
> Adjoint for matrix exp is fairly straightforward and indeed it is [implemented in ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl/blob/ae37562ea1f16816a0d8fff24e0aca6cd594a40f/src/rulesets/LinearAlgebra/matfun.jl#L117)
> 
> ReverseDiff and ForwardDiff however won’t use ChainRules. They both need to be told how to differentiate that function. I didn’t check this but I’m guessing there is some non-Julia in the back that prevents them from doing it out of the box. ReverseDiff.jl [can be told to use Chainrules.](https://github.com/JuliaDiff/ReverseDiff.jl/blob/b669294586b0789ab10bbda223543bcec64cf8aa/src/macros.jl#L299) Maybe that’s the best option.

I see, thank you very much. I gave this a shot, but still couldn’t get it to work:

```julia
using ReverseDiff
M = rand(2,2)

ReverseDiff.@grad_from_chainrules exp(x::ReverseDiff.TrackedArray)
ReverseDiff.@grad_from_chainrules exp(x::ReverseDiff.TrackedReal)
ReverseDiff.gradient(x -> sum(exp(x)),M) #MethodError: no method matching iterate(::Nothing)

```

Or did I misunderstand how the macro is suposed to be used?

> [@cortner](#):
>
> I’m more concerned about the type instability when you use Zygote. I can reproduce this. Maybe somebody should comment whether this is a Zygote or a ChainRules issue. I’m guessing it is Zygote, since calling `rrule` directly is type stable

I would guess it’s Zygote: I’ve had plenty of similar issues in the process of building my application, which is why refactored everything to make use of ReverseDiff (and would prefer not to use Zygote).

> [@cortner](#):
>
> P.S.: I’m assuming your example is just a MWE. Otherwise I would say just implement it yourself, possibly using `ChainRules.rrule`.

Yeah, it’s supposed to be part of a custom layer in a Neural Network.

---

<div class="post-metadata">

**Author:** ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)\
**Post date:** [December 21, 2023, 2:23pm UTC](https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865/5 "2023-12-21T14:23:45Z")

</div>

First off: great blog post, learned a lot from it.  
Still, given that my actual application will be in the context of neural networks, I think reverse-mode differentiation would be a better fit. Thank you, in any case.

---

<div class="post-metadata">

**Author:** ![facusapienza](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/facusapienza/32/24317_2.png) [@facusapienza](https://discourse.julialang.org/u/facusapienza)\
**Post date:** [December 21, 2023, 5:28pm UTC](https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865/6 "2023-12-21T17:28:42Z")

</div>

My guess is that the `exp` function primitive is not defined for exponential of a matrix? If instead you used a hand-written version of the exponential function (which approximates quite well the actual exponential)

```julia
function exp2(m)
    res = zeros(size(m))
    for i in 0:10
        res += m^i ./ factorial(i)
    end
    res
end

```

this works when doing the reverse pass

```julia
ReverseDiff.gradient(x -> sum(exp2(x)),M)

```

You can [customized pullback rules in ChainRulesBase](https://juliadiff.org/ChainRulesCore.jl/v1.11/rule_author/writing_good_rules.html#Constructors) and add them to your set of rules. Something like this should do the job, although I didn’t manage to make it work

```julia
function ChainRulesCore.rrule(::typeof(exp), M::AbstractMatrix)
    exp_pullback(ΔM) = (NoTangent(), exp(M) * ΔM)
    return exp(M), exp_pullback
end

```

Hope this helps. Please post once you find a solution to this! I would like to know how to make this work too 😉

---

<div class="post-metadata">

**Author:** ![cortner](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cortner/32/204_2.png) [@cortner](https://discourse.julialang.org/u/cortner)\
**Post date:** [December 21, 2023, 7:49pm UTC](https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865/7 "2023-12-21T19:49:22Z")

</div>

It actually is defined, see the posts above.

> <https://github.com/JuliaDiff/ChainRules.jl/blob/ae37562ea1f16816a0d8fff24e0aca6cd594a40f/src/rulesets/LinearAlgebra/matfun.jl#L117>

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [December 21, 2023, 8:05pm UTC](https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865/8 "2023-12-21T20:05:48Z")

</div>

> [@facusapienza](#):
>
> If instead you used a hand-written version of the exponential function (which approximates quite well the actual exponential)

This Taylor-series algorithm is discussed in section 3 (“method 1”) of the classic paper [Nineteen dubious ways to compute the exponential of a matrix](https://www.math.purdue.edu/~yipn/543/matrixExp19-I.pdf) (1978). Basically, the naive series is quite unreliable, even if you sum enough terms, because it is susceptible to [catastrophic cancellation](https://en.wikipedia.org/wiki/Catastrophic_cancellation).

(Contrary to popular misconception from first-year calculus, Taylor series are _not_ typically how special functions are computed.)

It’s really much safer to use the built-in `exp` function here, which means that you need to teach your AD system to use a custom rule (e.g. the one from ChainRules.jl).

---

<div class="post-metadata">

**Author:** ![cortner](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cortner/32/204_2.png) [@cortner](https://discourse.julialang.org/u/cortner)\
**Post date:** [December 21, 2023, 8:10pm UTC](https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865/9 "2023-12-21T20:10:50Z")

</div>

The following works fine for me:

```julia

using ReverseDiff, LinearAlgebra, ChainRules, Zygote
ReverseDiff.@grad_from_chainrules LinearAlgebra.exp(x::ReverseDiff.TrackedArray)
A = randn(3,3)
F = x -> sum(exp(x))
G1 = ReverseDiff.gradient(F, A)
G2 = Zygote.gradient(F, A)[1]
@show G1 ≈ G2

```

P.S.: @Bizzi – I’m not sure why your script didn’t work. Probably some combination of imports/using that are required to run this. This is where I sometimes get confused and just try until it runs ☹

P.P.S.: neither Zygote nor ReverseDiff have great performance here…

```julia
using BenchmarkTools
@btime ReverseDiff.gradient($F, $A) 
# 5.611 μs (81 allocations: 8.34 KiB)
@btime Zygote.gradient($F, $A)
# 3.970 μs (74 allocations: 8.03 KiB)

```

---

<div class="post-metadata">

**Author:** ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)\
**Post date:** [December 22, 2023, 4:27pm UTC](https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865/10 "2023-12-22T16:27:29Z")

</div>

You’re right, importing `ChainRules` was also necessary 🫠. Kinda misleading error message but anyway.  
Thank you all very much!

---

<div class="post-metadata">

**Author:** ![mikmoore](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mikmoore/32/31109_2.png) [@mikmoore](https://discourse.julialang.org/u/mikmoore)\
**Post date:** [January 2, 2024, 3:48pm UTC](https://discourse.julialang.org/t/taking-gradients-of-a-matrix-exponential/107865/11 "2024-01-02T15:48:45Z")

</div>

Related issue [#51008](https://github.com/JuliaLang/julia/issues/51008). `exp` just doesn’t work with many input types (it only works for matrices with `eltype`s that promote to `LinearAlgebra.BlasFloat`, which does not include things like dual numbers). Fixing it doesn’t require any significant algorithmic adjustments, just that someone go through and add a code path for more generic types.