Fuse frule

Author: YingboMa

Project.toml

@@ -10,7 +10,9 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-ChainRulesCore = "0.4"
+Reexport = "0.2"
+Requires = "0.5.2"
+ChainRulesCore = "0.5"
 FiniteDifferences = "^0.7"
 Reexport = "0.2"
 Requires = "0.5.2, 1"

src/rulesets/Base/base.jl

@@ -100,11 +100,8 @@
 
 # product rule requires special care for arguments where `mul` is non-commutative
 
-function frule(::typeof(*), x::Number, y::Number)
-    function times_pushforward(_, Δx, Δy)
-        return Δx * y + x * Δy
-    end
-    return x * y, times_pushforward
+function frule(::typeof(*), x::Number, y::Number, _, Δx, Δy)
+    return x * y, Δx * y + x * Δy
 end
 
 function rrule(::typeof(*), x::Number, y::Number)
@@ -114,11 +111,8 @@ function rrule(::typeof(*), x::Number, y::Number)
     return x * y, times_pullback
 end
 
-function frule(::typeof(identity), x)
-    function identity_pushforward(_, ẏ)
-        return ẏ
-    end
-    return x, identity_pushforward
+function frule(::typeof(identity), x, _, ẏ)
+    return x, ẏ
 end
 
 function rrule(::typeof(identity), x)

src/rulesets/Base/broadcast.jl

@@ -6,19 +6,17 @@ https://github.com/JuliaLang/julia/issues/22129.
 =#
 function _cast_diff(f, x)
     function element_rule(u)
-        fu, du = frule(f, u)
-        fu, extern(du(NamedTuple(), One()))
+        dself = Zero()
+        fu, du = frule(f, u, dself, One())
+        fu, extern(du)
     end
     results = broadcast(element_rule, x)
     return first.(results), last.(results)
 end
 
-function frule(::typeof(broadcast), f, x)
+function frule(::typeof(broadcast), f, x, _, Δf, Δx)
     Ω, ∂x = _cast_diff(f, x)
-    function broadcast_pushforward(_, Δf, Δx)
-        return Δx .* ∂x
-    end
-    return Ω, broadcast_pushforward
+    return Ω, Δx .* ∂x
 end
 
 function rrule(::typeof(broadcast), f, x)

src/rulesets/Base/mapreduce.jl

@@ -50,11 +50,8 @@ end
 ##### `sum`
 #####
 
-function frule(::typeof(sum), x)
-    function sum_pushforward(_, ẋ)
-        return sum(ẋ)
-    end
-    return sum(x), sum_pushforward
+function frule(::typeof(sum), x, _, ẋ)
+    return sum(x), sum(ẋ)
 end
 
 function rrule(::typeof(sum), f, x::AbstractArray{<:Real}; dims=:)

src/rulesets/LinearAlgebra/blas.jl

@@ -11,7 +11,7 @@ _zeros(x) = fill!(similar(x), zero(eltype(x)))
 ##### `BLAS.dot`
 #####
 
-frule(::typeof(BLAS.dot), x, y) = frule(dot, x, y)
+frule(::typeof(BLAS.dot), x, y, Δself, Δx, Δy) = frule(dot, x, y, Δself, Δx, Δy)
 
 rrule(::typeof(BLAS.dot), x, y) = rrule(dot, x, y)
 
@@ -35,12 +35,9 @@ end
 ##### `BLAS.nrm2`
 #####
 
-function frule(::typeof(BLAS.nrm2), x)
+function frule(::typeof(BLAS.nrm2), x, _, Δ)
     Ω = BLAS.nrm2(x)
-    function nrm2_pushforward(_, Δx)
-        return sum(Δx * cast(@thunk(x * inv(Ω))))
-    end
-    return Ω, nrm2_pushforward
+    return Ω, sum(Δx * cast(@thunk(x * inv(Ω))))
 end
 
 function rrule(::typeof(BLAS.nrm2), x)
@@ -70,11 +67,8 @@ end
 ##### `BLAS.asum`
 #####
 
-function frule(::typeof(BLAS.asum), x)
-    function asum_pushforward(_, Δx)
-        return sum(cast(sign, x) * Δx)
-    end
-    return BLAS.asum(x), asum_pushforward
+function frule(::typeof(BLAS.asum), x, _, Δx)
+    return BLAS.asum(x), sum(cast(sign, x) * Δx)
 end
 
 function rrule(::typeof(BLAS.asum), x)

src/rulesets/LinearAlgebra/dense.jl

@@ -8,11 +8,8 @@ const SquareMatrix{T} = Union{Diagonal{T},AbstractTriangular{T}}
 ##### `dot`
 #####
 
-function frule(::typeof(dot), x, y)
-    function dot_pushforward(Δself, Δx, Δy)
-        return sum(Δx .* y) + sum(x .* Δy)
-    end
-    return dot(x, y), dot_pushforward
+function frule(::typeof(dot), x, y, _, Δx, Δy)
+    return dot(x, y), sum(Δx .* y) + sum(x .* Δy)
 end
 
 function rrule(::typeof(dot), x, y)
@@ -26,13 +23,10 @@ end
 ##### `inv`
 #####
 
-function frule(::typeof(inv), x::AbstractArray)
+function frule(::typeof(inv), x::AbstractArray, _, Δx)
     Ω = inv(x)
     m = @thunk(-Ω)
-    function inv_pushforward(_, Δx)
-        return m * Δx * Ω
-    end
-    return Ω, inv_pushforward
+    return Ω, m * Δx * Ω
 end
 
 function rrule(::typeof(inv), x::AbstractArray)
@@ -48,14 +42,11 @@ end
 ##### `det`
 #####
 
-function frule(::typeof(det), x)
+function frule(::typeof(det), x, _, ẋ)
     Ω = det(x)
-    function det_pushforward(_, ẋ)
-        # TODO Performance optimization: probably there is an efficent
-        # way to compute this trace without during the full compution within
-        return Ω * tr(inv(x) * ẋ)
-    end
-    return Ω, det_pushforward
+    # TODO Performance optimization: probably there is an efficent
+    # way to compute this trace without during the full compution within
+    return Ω, Ω * tr(inv(x) * ẋ)
 end
 
 function rrule(::typeof(det), x)
@@ -70,12 +61,9 @@ end
 ##### `logdet`
 #####
 
-function frule(::typeof(logdet), x)
+function frule(::typeof(logdet), x, _, Δx)
     Ω = logdet(x)
-    function logdet_pushforward(_, Δx)
-        return tr(inv(x) * Δx)
-    end
-    return Ω, logdet_pushforward
+    return Ω, tr(inv(x) * Δx)
 end
 
 function rrule(::typeof(logdet), x)
@@ -90,11 +78,8 @@ end
 ##### `trace`
 #####
 
-function frule(::typeof(tr), x)
-    function tr_pushforward(_, Δx)
-        return tr(Δx)
-    end
-    return tr(x), tr_pushforward
+function frule(::typeof(tr), x, _, Δx)
+    return tr(x), tr(Δx)
 end
 
 function rrule(::typeof(tr), x)

test/rulesets/Base/base.jl

@@ -51,13 +51,32 @@
             test_scalar(acscd, 1/x)
             test_scalar(acotd, 1/x)
         end
-        
-        @testset "sincos" begin
-            x, Δx, x̄ = randn(3)
-            Δz = (randn(), randn())
-
-            frule_test(sincos, (x, Δx))
-            rrule_test(sincos, Δz, (x, x̄))
+        @testset "Multivariate" begin
+            @testset "atan2" begin
+                # https://en.wikipedia.org/wiki/Atan2
+                x, y = rand(2)
+                ratan = atan(x, y)
+                u = x^2 + y^2
+                datan = y/u - 2x/u
+
+                r, ṙ = frule(atan, x, y, Zero(), 1, 2)
+                @test r === ratan
+                @test ṙ === datan
+
+                r, pullback = rrule(atan, x, y)
+                @test r === ratan
+                dself, df1, df2 = pullback(1)
+                @test dself == NO_FIELDS
+                @test df1 + 2df2 === datan
+            end
+
+            @testset "sincos" begin
+                x, Δx, x̄ = randn(3)
+                Δz = (randn(), randn())
+
+                frule_test(sincos, (x, Δx))
+                rrule_test(sincos, Δz, (x, x̄))
+            end
         end
     end  # Trig
 
@@ -128,8 +147,7 @@
             _, x̄ = pb(10.5)
             @test extern(x̄) == 0
 
-            _, pf = frule(sign, 0.0)
-            ẏ = pf(NamedTuple(), 10.5)
+            _, ẏ = frule(sign, 0.0, Zero(), 10.5)
             @test extern(ẏ) == 0
         end
     end

test/rulesets/Base/broadcast.jl

@@ -22,10 +22,8 @@
         end
         @testset "frule" begin
             x = rand(3, 3)
-            y, pushforward = frule(broadcast, sin, x)
+            y, ẏ = frule(broadcast, sin, x, Zero(), Zero(), One())
             @test y == sin.(x)
-
-            ẏ = pushforward(NamedTuple(), NamedTuple(), One())
             @test extern(ẏ) == cos.(x)
         end
     end

test/test_util.jl

@@ -69,7 +69,6 @@ function frule_test(f, xẋs::Tuple{Any, Any}...; rtol=1e-9, atol=1e-9, fdm=_fdm
     Ω, dΩ_ad = ChainRules.frule(f, xs..., dself, ẋs...)
     dΩ_ad = dΩ_ad .+ 0
     @test f(xs...) == Ω
-    dΩ_ad = pushforward(NamedTuple(), ẋs...)
 
     # Correctness testing via finite differencing.
     dΩ_fd = jvp(fdm, xs->f(xs...), (xs, ẋs))