Fix #211 making 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

test/rule_definition_tools.jl

@@ -110,4 +110,20 @@ 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 with === because type also must match
+            @test ẏ === Composite{typeof(y)}(50f0, 100f0)
+        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