V0.2.1 (#18)

* fix #14 (#16) * fix #14 * add FD * Update test/runtests.jl Co-authored-by: Mohamed Tarek <[email protected]> --------- Co-authored-by: Mohamed Tarek <[email protected]> Co-authored-by: ThummeTo <[email protected]> * version inc * added `FiniteDifferences` --------- Co-authored-by: Fredrik Bagge Carlson <[email protected]> Co-authored-by: Mohamed Tarek <[email protected]>
ThummeTo · Sep 14, 2023 · ef8b579 · ef8b579 · ThummeTo · Sep 14, 2023
1 parent 978c246
commit ef8b579
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 7 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "ForwardDiffChainRules"
 uuid = "c9556dd2-1aed-4cfe-8560-1557cf593001"
 authors = ["MT <[email protected]>", "TT <[email protected]>"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -10,7 +10,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 
 [compat]
-ChainRulesCore = "1.15.0"
+ChainRulesCore = "1.15.0, 1.16.0"
 DifferentiableFlatten = "0.1.1"
 ForwardDiff = "0.10.0"
 MacroTools = "0.5.10"

diff --git a/src/ForwardDiffChainRules.jl b/src/ForwardDiffChainRules.jl
@@ -35,8 +35,19 @@ import ForwardDiff
 import MacroTools
 using DifferentiableFlatten
 
-macro ForwardDiff_frule(sig)
-    _fd_frule(sig)
+"""
+    @ForwardDiff_frule signature mutating
+
+`mutating` indicates whether or not the function is mutating the input argument.
+
+# Example
+To define a rule for `LinearAlgebra.exp!`, which is a mutating funciton, we call the macro like this
+```julia
+@ForwardDiff_frule LinearAlgebra.exp!(A::AbstractMatrix{<:ForwardDiff.Dual}) true
+```
+"""
+macro ForwardDiff_frule(sig, mutating=false)
+    _fd_frule(sig; mutating)
 end
 export @ForwardDiff_frule
 
@@ -59,7 +70,7 @@ end
 
 const cfg = ForwardDiffRuleConfig()
 
-function _fd_frule(sig)
+function _fd_frule(sig; mutating=false)
     if MacroTools.@capture(sig, f_(x__; k__))
         nothing
     else
@@ -77,15 +88,21 @@ function _fd_frule(sig)
             flat_xpartials = reduce(vcat, transpose.(ForwardDiff.partials.(flatx)))
 
             xprimals = unflattenx(flat_xprimals)
+            xprimals_copy = $mutating ? copy.(xprimals) : xprimals
             xpartials1 = unflattenx(flat_xpartials[:,1])
             yprimals, ypartials1 = ChainRulesCore.frule(
-                cfg, (NoTangent(), xpartials1...), f, xprimals...; ks...,
+                cfg, (NoTangent(), xpartials1...), f, xprimals_copy...; ks...,
             )
             flat_yprimals, unflatteny = ForwardDiffChainRules.DifferentiableFlatten.flatten(yprimals)
             flat_ypartials1, _ = ForwardDiffChainRules.DifferentiableFlatten.flatten(ypartials1)
             flat_ypartials = hcat(reshape(flat_ypartials1, :, 1), ntuple(Val(CS - 1)) do i
+                if $mutating
+                    for (xpc, xp) in zip(xprimals_copy, xprimals)
+                        xpc .= xp # Update copy
+                    end
+                end
                 xpartialsi = unflattenx(flat_xpartials[:, i+1])
-                _, ypartialsi = ChainRulesCore.frule(cfg, (NoTangent(), xpartialsi...), f, xprimals...; ks...)
+                _, ypartialsi = ChainRulesCore.frule(cfg, (NoTangent(), xpartialsi...), f, xprimals_copy...; ks...)
                 return ForwardDiffChainRules.DifferentiableFlatten.flatten(ypartialsi)[1]
             end...)
 

diff --git a/test/Project.toml b/test/Project.toml
@@ -1,6 +1,7 @@
 [deps]
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NamedTupleTools = "d9ec5142-1e00-5aa0-9d6a-321866360f50"

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -9,6 +9,7 @@ using Test
 using ChainRules, ChainRulesCore, ForwardDiff
 using LinearAlgebra
 import NamedTupleTools: ntfromstruct, structfromnt
+import FiniteDifferences
 
 #
 # Copyright (c) 2022 The contributors
@@ -56,6 +57,7 @@ SOFTWARE.
         @ForwardDiff_frule f1(x1::ForwardDiff.Dual, x2::ForwardDiff.Dual)
         @ForwardDiff_frule f1(x1::AbstractVector{<:ForwardDiff.Dual}, x2::AbstractVector{<:ForwardDiff.Dual})
         @ForwardDiff_frule f1(x1::AbstractMatrix{<:ForwardDiff.Dual}, x2::AbstractMatrix{<:ForwardDiff.Dual})
+        @ForwardDiff_frule LinearAlgebra.exp!(A::AbstractMatrix{<:ForwardDiff.Dual}) true
 
         f2(x::NamedTuple, y::NamedTuple) = (a = x.a + y.a, b = x.b + y.b)
         f2(x::AbstractVector, y::AbstractVector) = f2.(x, y)
@@ -168,6 +170,13 @@ SOFTWARE.
         @test frule_count == 16
         @test norm(g - I) < 1e-6
     end
+    @testset "exp!" begin
+        fexp = x -> sum(LinearAlgebra.exp!(copy(x)))
+        X = rand(4, 4)
+        g = ForwardDiff.gradient(fexp, X)
+        g2 = FiniteDifferences.grad(FiniteDifferences.central_fdm(5, 1), fexp, X)[1]
+        @test norm(g-g2) < 1e-4
+    end
     @testset "kwargs" begin
         fkwarg(x1, x2; a = 2.0) = x1 * x2 * a
         frule_count = 0