Merge pull request #214 from JuliaDiff/ox/scalarmultiple

Make scalar_rule's frule return a Composite not a tuple

Author: oxinabox

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.9.12"
+version = "0.9.13"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/rule_definition_tools.jl

@@ -148,10 +148,12 @@ function scalar_frule_expr(f, call, setup_stmts, inputs, partials)
         propagation_expr(Δs, ∂s)
     end
     if n_outputs > 1
-        # For forward-mode we only return a tuple if output actually a tuple.
-        pushforward_returns = Expr(:tuple, pushforward_returns...)
+        # For forward-mode we return a Composite if output actually a tuple.
+        pushforward_returns = Expr(
+            :call, :(ChainRulesCore.Composite{typeof($(esc(:Ω)))}), pushforward_returns...
+        )
     else
-        pushforward_returns = pushforward_returns[1]
+        pushforward_returns = first(pushforward_returns)
     end
 
     return quote

src/rules.jl

@@ -29,18 +29,30 @@ julia> Δsinx == cos(x)
 true
 ```
 
-unary input, binary output scalar function:
+Unary input, binary output scalar function:
 
 ```jldoctest frule
 julia> sincosx, Δsincosx = frule((dself, 1), sincos, x);
 
 julia> sincosx == sincos(x)
 true
 
-julia> Δsincosx == (cos(x), -sin(x))
+julia> Δsincosx[1] == cos(x)
+true
+
+julia> Δsincosx[2] == -sin(x)
 true
 ```
 
+Note that techically speaking julia does not have multiple output functions, just functions
+that return a single output that is iterable, like a `Tuple`.
+So this is actually a [`Composite`](@ref):
+```jldoctest frule
+julia> Δsincosx
+Composite{Tuple{Float64,Float64}}(0.6795498147167869, -0.7336293678134624)
+```.
+
+
 See also: [`rrule`](@ref), [`@scalar_rule`](@ref)
 """
 frule(::Any, ::Vararg{Any}; kwargs...) = nothing

test/rule_definition_tools.jl

@@ -110,4 +110,21 @@ end
             )
         end
     end
+
+    @testset "@scalar_rule" begin
+        @testset "@scalar_rule with multiple output" begin
+            simo(x) = (x, 2x)
+            @scalar_rule(simo(x), 1f0, 2f0)
+    
+            y, simo_pb = rrule(simo, π)
+
+            @test simo_pb((10f0, 20f0)) == (NO_FIELDS, 50f0)
+
+            y, ẏ = frule((NO_FIELDS, 50f0), simo, π)
+            @test y == (π, 2π)
+            @test ẏ == Composite{typeof(y)}(50f0, 100f0)
+            # make sure type is exactly as expected:
+            @test ẏ isa Composite{Tuple{Irrational{:π}, Float64}, Tuple{Float32, Float32}}
+        end
+    end
 end

test/rules.jl

@@ -136,18 +136,4 @@ _second(t) = Base.tuple_type_head(Base.tuple_type_tail(t))
         @test ∂self == NO_FIELDS
         @test ∂x ≈ j′vp(central_fdm(5, 1), complex_times, Ω̄, x)[1]
     end
-end
-
-
-simo(x) = (x, 2x)
-@scalar_rule(simo(x), 1, 2)
-
-@testset "@scalar_rule with multiple inputs" begin
-    y, simo_pb = rrule(simo, π)
-
-    @test simo_pb((10, 20)) == (NO_FIELDS, 50)
-
-    y, ẏ = frule((NO_FIELDS, 50), simo, π)
-    @test y == (π, 2π)
-    @test ẏ == (50, 100)
-end
+end