Merge pull request #88 from JuliaDiff/myb/fuse_frule

Fuse `frule` and pushforward

Author: YingboMa

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.5.0-DEV"
+version = "0.5.0"
 
 [compat]
 julia = "^1.0"

src/differentials/composite.jl

@@ -86,7 +86,7 @@ function backing(x::T)::NamedTuple where T
         nfields = fieldcount(T)
         names = fieldnames(T)
         types = fieldtypes(T)
-        
+
         vals = Expr(:tuple, ntuple(ii->:(getfield(x, $ii)), nfields)...)
         return :(NamedTuple{$names, Tuple{$(types...)}}($vals))
     else

src/differentials/one.jl

@@ -11,4 +11,3 @@ Base.Broadcast.broadcastable(::One) = Ref(One())
 
 Base.iterate(x::One) = (x, nothing)
 Base.iterate(::One, ::Any) = nothing
-

src/differentials/zero.jl

@@ -11,5 +11,3 @@ Base.Broadcast.broadcastable(::Zero) = Ref(Zero())
 
 Base.iterate(x::Zero) = (x, nothing)
 Base.iterate(::Zero, ::Any) = nothing
-
-

src/rule_definition_tools.jl

@@ -159,19 +159,13 @@ function scalar_frule_expr(f, call, setup_stmts, inputs, partials)
         pushforward_returns = pushforward_returns[1]
     end
 
-    pushforward = quote
+    return quote
         # _ is the input derivative w.r.t. function internals. since we do not
         # allow closures/functors with @scalar_rule, it is always ignored
-        function $(propagator_name(f, :pushforward))(_, $(Δs...))
-            $pushforward_returns
-        end
-    end
-
-    return quote
-        function ChainRulesCore.frule(::typeof($f), $(inputs...))
+        function ChainRulesCore.frule(::typeof($f), $(inputs...), _, $(Δs...))
             $(esc(:Ω)) = $call
             $(setup_stmts...)
-            return $(esc(:Ω)), $pushforward
+            return $(esc(:Ω)), $pushforward_returns
         end
     end
 end

src/rules.jl

@@ -2,34 +2,40 @@
 ##### `frule`/`rrule`
 #####
 
+# TODO: remember to update the examples
 """
-    frule(f, x...)
+    frule(f, x..., ṡelf, Δx...)
 
-Expressing `x` as the tuple `(x₁, x₂, ...)` and the output tuple of `f(x...)`
-as `Ω`, return the tuple:
+Expressing `x` as the tuple `(x₁, x₂, ...)`, `Δx` as the tuple `(Δx₁, Δx₂,
+...)`, and the output tuple of `f(x...)` as `Ω`, return the tuple:
 
-    (Ω, (ṡelf, ẋ₁, ẋ₂, ...) -> Ω̇₁, Ω̇₂, ...)
+    (Ω, (Ω̇₁, Ω̇₂, ...))
 
 The second return value is the propagation rule, or the pushforward.
 It takes in differentials corresponding to the inputs (`ẋ₁, ẋ₂, ...`)
 and `ṡelf` the internal values of the function (for closures).
 
 
-If no method matching `frule(f, xs...)` has been defined, then return `nothing`.
+If no method matching `frule(f, x..., ṡelf, Δx...)` has been defined, then
+return `nothing`.
 
 Examples:
 
 unary input, unary output scalar function:
 
 ```
+julia> dself = Zero()
+Zero()
+
 julia> x = rand();
 
-julia> sinx, sin_pushforward = frule(sin, x);
+julia> sinx, sin_pushforward = frule(sin, x, dself, 1)
+(0.35696518021277485, 0.9341176907197836)
 
 julia> sinx == sin(x)
 true
 
-julia> sin_pushforward(NamedTuple(), 1) == cos(x)
+julia> sin_pushforward == cos(x)
 true
 ```
 
@@ -38,12 +44,12 @@ unary input, binary output scalar function:
 ```
 julia> x = rand();
 
-julia> sincosx, sincos_pushforward = frule(sincos, x);
+julia> sincosx, sincos_pushforward = frule(sincos, x, dself, 1);
 
 julia> sincosx == sincos(x)
 true
 
-julia> sincos_pushforward(NamedTuple(), 1) == (cos(x), -sin(x))
+julia> sincos_pushforward == (cos(x), -sin(x))
 true
 ```
 

test/rules.jl

@@ -13,8 +13,9 @@ dummy_identity(x) = x
 _second(t) = Base.tuple_type_head(Base.tuple_type_tail(t))
 
 @testset "frule and rrule" begin
-    @test frule(cool, 1) === nothing
-    @test frule(cool, 1; iscool=true) === nothing
+    dself = Zero()
+    @test frule(cool, 1, dself, 1) === nothing
+    @test frule(cool, 1, dself, 1; iscool=true) === nothing
     @test rrule(cool, 1) === nothing
     @test rrule(cool, 1; iscool=true) === nothing
 
@@ -29,9 +30,9 @@ _second(t) = Base.tuple_type_head(Base.tuple_type_tail(t))
                         Tuple{typeof(rrule),typeof(cool),String}])
     @test cool_methods == only_methods
 
-    frx, cool_pushforward = frule(cool, 1)
+    frx, cool_pushforward = frule(cool, 1, dself, 1)
     @test frx == 2
-    @test cool_pushforward(NamedTuple(), 1) == 1
+    @test cool_pushforward == 1
     rrx, cool_pullback = rrule(cool, 1)
     self, rr1 = cool_pullback(1)
     @test self == NO_FIELDS