Import matrix exponential rule from ChainRules (#365)

* Matrix exponential

* Bump patch version

* Remove bad rule definition

Project.toml

@@ -1,10 +1,11 @@
 name = "Mooncake"
 uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.4.40"
+version = "0.4.41"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
@@ -43,6 +44,7 @@ AllocCheck = "0.2"
 Aqua = "0.8.9"
 BenchmarkTools = "1"
 CUDA = "5"
+ChainRules = "1.71.0"
 ChainRulesCore = "1"
 DiffRules = "1"
 DiffTests = "0.1"
@@ -56,8 +58,8 @@ LogDensityProblemsAD = "1"
 LuxLib = "1.2 - 1.3.3"
 MistyClosures = "2"
 NNlib = "0.9"
-Random = "1"
 Pkg = "1"
+Random = "1"
 Setfield = "1"
 SpecialFunctions = "2"
 StableRNGs = "1"

src/Mooncake.jl

@@ -4,6 +4,7 @@ const CC = Core.Compiler
 
 using
     ADTypes,
+    ChainRules,
     DiffRules,
     ExprTools,
     Graphs,
@@ -83,6 +84,7 @@ include(joinpath("rrules", "fastmath.jl"))
 include(joinpath("rrules", "foreigncall.jl"))
 include(joinpath("rrules", "iddict.jl"))
 include(joinpath("rrules", "lapack.jl"))
+include(joinpath("rrules", "linear_algebra.jl"))
 include(joinpath("rrules", "low_level_maths.jl"))
 include(joinpath("rrules", "misc.jl"))
 include(joinpath("rrules", "new.jl"))

src/rrules/linear_algebra.jl

@@ -0,0 +1,32 @@
+@is_primitive MinimalCtx Tuple{typeof(exp), Matrix{<:IEEEFloat}}
+
+struct ExpPullback{P}
+    pb
+    Ȳ::Matrix{P}
+    X̄::Matrix{P}
+end
+
+function (pb::ExpPullback)(::NoRData)
+    _, X̄_inc = pb.pb(pb.Ȳ)
+    pb.X̄ .+= X̄_inc
+    return NoRData(), NoRData()
+end
+
+function rrule!!(::CoDual{typeof(exp)}, X::CoDual{Matrix{P}}) where {P<:IEEEFloat}
+    Y, pb = ChainRules.rrule(exp, X.x)
+    Ȳ = zero(Y)
+    return CoDual(Y, Ȳ), ExpPullback{P}(pb, Ȳ, X.dx)
+end
+
+function generate_hand_written_rrule!!_test_cases(rng_ctor, ::Val{:linear_algebra})
+    test_cases = Any[
+        (false, :none, nothing, exp, randn(3, 3)),
+        (false, :none, nothing, exp, randn(7, 7)),
+    ]
+    memory = Any[]
+    return test_cases, memory
+end
+
+function generate_derived_rrule!!_test_cases(rng_ctor, ::Val{:linear_algebra})
+    return Any[], Any[]
+end

test/rrules/linear_algebra.jl

@@ -0,0 +1,3 @@
+@testset "linear_algebra" begin
+    TestUtils.run_rrule!!_test_cases(StableRNG, Val(:linear_algebra))
+end

test/runtests.jl

@@ -40,6 +40,8 @@ include("front_matter.jl")
             include(joinpath("rrules", "iddict.jl"))
             @info "lapack"
             include(joinpath("rrules", "lapack.jl"))
+            @info "linear_algebra"
+            include(joinpath("rrules", "linear_algebra.jl"))
             @info "low_level_maths"
             include(joinpath("rrules", "low_level_maths.jl"))
             @info "misc"