diff --git a/Project.toml b/Project.toml index 1fdcbad09..c80a71522 100644 --- a/Project.toml +++ b/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" diff --git a/src/Mooncake.jl b/src/Mooncake.jl index 6889da23e..0b5a222ac 100644 --- a/src/Mooncake.jl +++ b/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")) diff --git a/src/rrules/linear_algebra.jl b/src/rrules/linear_algebra.jl new file mode 100644 index 000000000..dc628d2bc --- /dev/null +++ b/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 diff --git a/test/rrules/linear_algebra.jl b/test/rrules/linear_algebra.jl new file mode 100644 index 000000000..f7e1ac0ed --- /dev/null +++ b/test/rrules/linear_algebra.jl @@ -0,0 +1,3 @@ +@testset "linear_algebra" begin + TestUtils.run_rrule!!_test_cases(StableRNG, Val(:linear_algebra)) +end diff --git a/test/runtests.jl b/test/runtests.jl index dd9d955ea..a4a9e7d71 100644 --- a/test/runtests.jl +++ b/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"