Merge pull request #36 from JuliaDiff/ox/1pb

Remove special handling for pullbacks having multiple inputs

Author: oxinabox

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesTestUtils"
 uuid = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-version = "0.2.7"
+version = "0.3"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -11,7 +11,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ChainRulesCore = "0.7.1"
+ChainRulesCore = "0.8"
 Compat = "3"
 FiniteDifferences = "0.9"
 julia = "1"

src/testers.jl

@@ -147,9 +147,8 @@ function rrule_test(f, ȳ, xx̄s::Tuple{Any, Any}...; rtol=1e-9, atol=1e-9, fdm
     # use collect so can do vector equality
     @test isapprox(collect(y_ad), collect(y); rtol=rtol, atol=atol)
     @assert !(isa(ȳ, Thunk))
-    # If the function returned multiple values,
-    # then it must have multiple seeds for propagating backwards
-    ∂s = (y_ad isa Tuple) ? pullback(ȳ...) : pullback(ȳ)
+    
+    ∂s = pullback(ȳ)
     ∂self = ∂s[1]
     x̄s_ad = ∂s[2:end]
     @test ∂self === NO_FIELDS  # No internal fields

test/testers.jl

@@ -49,6 +49,28 @@ fbtestkws(x, y; err = true) = err ? error() : x
         end
     end
 
+    @testset "single input, multiple output" begin
+        simo(x) = (x, 2x)
+        function ChainRulesCore.rrule(simo, x)
+            simo_pullback((a, b)) = (NO_FIELDS, a .+ 2 .* b)
+            return simo(x), simo_pullback
+        end
+        function ChainRulesCore.frule((_, ẋ), simo, x)
+            y = simo(x)
+            return y, Composite{typeof(y)}(ẋ, 2ẋ)
+        end
+
+        @testset "frule_test" begin
+            frule_test(simo, (randn(), randn()))  # on scalar
+            frule_test(simo, (randn(4), randn(4)))  # on array
+        end
+        @testset "rrule_test" begin
+            # note: we are pulling back tuples (could use Composites here instead)
+            rrule_test(simo, (randn(), rand()), (randn(), randn()))  # on scalar
+            rrule_test(simo, (randn(4), rand(4)), (randn(4), randn(4))) # on array
+        end
+    end
+
 
     @testset "tuple input: first" begin
         ChainRulesCore.frule((_, dx), ::typeof(first), xs::Tuple) = (first(xs), first(dx))