fix tests

Author: mcabbott

src/rulesets/Base/broadcast.jl

@@ -240,7 +240,7 @@ end
 rrule(::RCR, ::typeof(broadcasted), ::typeof(+), args::Number...) = rrule(+, args...) |> _prepend_zero
 rrule(::RCR, ::typeof(broadcasted), ::typeof(-), x::Number, y::Number) = rrule(-, x, y) |> _prepend_zero
 rrule(::RCR, ::typeof(broadcasted), ::typeof(-), x::Number) = rrule(-, x) |> _prepend_zero
-rrule(::RCR, ::typeof(broadcasted), ::typeof(*), x::Number, y::Number) = rrule(*, x, y) |> _prepend_zero
+rrule(::RCR, ::typeof(broadcasted), ::typeof(*), args::Number...) = rrule(*, args...) |> _prepend_zero
 rrule(::RCR, ::typeof(broadcasted), ::typeof(Base.literal_pow), ::typeof(^), x::Number, ::Val{2}) =
     rrule(Base.literal_pow, ^, x, Val(2)) |> _prepend_zero
 rrule(::RCR, ::typeof(broadcasted), ::typeof(/), x::Number, y::Number) = rrule(/, x, y) |> _prepend_zero

src/tuplecast.jl

@@ -3,7 +3,7 @@
     tuplecast(f, args...)
 
 For a function `f` which returns a tuple, this is `== unzip(broadcast(f, args...))`,
-but performed using `StructArrays` for efficiency.
+but performed using `StructArrays` for efficiency. Used in the gradient of broadcasting.
 
 # Examples
 ```
@@ -52,7 +52,8 @@ function ChainRulesCore.rrule(cfg::RuleConfig{>:HasReverseMode}, ::typeof(tuplec
     return z, untuplecast
 end
 
-function rrule(cfg::RCR, ::typeof(collect∘tuplecast), f, args...)  # for testing, but doesn't work?
+# This is for testing, but the tests using it don't work.
+function rrule(cfg::RuleConfig{>:HasReverseMode}, ::typeof(collect∘tuplecast), f, args...)
     y, back = rrule(cfg, tuplecast, f, args...)
     return collect(y), back
 end
@@ -62,6 +63,8 @@ end
 
 For a function `f` which returns a tuple, this is `== unzip(map(f, args...))`,
 but performed using `StructArrays` for efficiency.
+
+Not in use at present, but see `tuplecast`.
 """
 function tuplemap(f::F, args...) where {F}
     T = Broadcast.combine_eltypes(f, args)

test/rulesets/Base/broadcast.jl

@@ -60,7 +60,7 @@ using Base.Broadcast: broadcasted
         test_rrule(copy∘broadcasted, /, (rand(2),), rand(3), check_inferred=false)
     end
 
-    @testset "lazy rules" begin
+    @testset "fused rules" begin
         @testset "arithmetic" begin
             test_rrule(copy∘broadcasted, +, rand(3), rand(3))
             test_rrule(copy∘broadcasted, +, rand(3), rand(4)')
@@ -88,8 +88,9 @@ using Base.Broadcast: broadcasted
             @test y4 == [im, 2im, 3im]
             @test unthunk(bk4([4, 5im, 6+7im])[4]) == [0,5,7]
 
-            test_rrule(copy∘broadcasted, *, rand(3), rand(3), rand(3), rand(3), rand(3))
-            test_rrule(copy∘broadcasted, *, rand(), rand(), rand(3), rand(3) .+ im, rand(4)')
+            test_rrule(copy∘broadcasted, *, rand(3), rand(3), rand(3), rand(3), rand(3), check_inferred=false)  # Union{NoTangent, ZeroTangent}
+            test_rrule(copy∘broadcasted, *, rand(), rand(), rand(3), rand(3) .+ im, rand(4)', check_inferred=false)  # Union{NoTangent, ZeroTangent}
+            # (These two may infer with vararg rrule)
 
             test_rrule(copy∘broadcasted, Base.literal_pow, ^, rand(3), Val(2))
             test_rrule(copy∘broadcasted, Base.literal_pow, ^, rand(3) .+ im, Val(2))
@@ -134,6 +135,7 @@ using Base.Broadcast: broadcasted
             test_rrule(copy∘broadcasted, -, rand(), rand())
             test_rrule(copy∘broadcasted, -, rand())
             test_rrule(copy∘broadcasted, *, rand(), rand())
+            test_rrule(copy∘broadcasted, *, rand(), rand(), rand(), rand())
             test_rrule(copy∘broadcasted, Base.literal_pow, ^, rand(), Val(2))
             test_rrule(copy∘broadcasted, /, rand(), rand())
         end

test/rulesets/Base/mapreduce.jl

@@ -2,8 +2,6 @@
 Base.sum(xs::AbstractArray, weights::AbstractArray) = dot(xs, weights)
 struct SumRuleConfig <: RuleConfig{Union{HasReverseMode}} end
 
-const CFG = ChainRulesTestUtils.ADviaRuleConfig()
-
 @testset "Reductions" begin
     @testset "sum(::Tuple)" begin
         test_frule(sum, Tuple(rand(5)))

test/runtests.jl

@@ -43,8 +43,7 @@ else
 end
 
 @testset "ChainRules" begin  # One overall @testset ensures it keeps going after failures
-    include("test_helpers.jl")
-    include("tuplecast.jl")
+    include("test_helpers.jl")  # This can't be skipped
     println()
 
     test_method_tables()  # Check the global method tables are consistent
@@ -60,6 +59,8 @@ end
     include_test("rulesets/Base/sort.jl")
     include_test("rulesets/Base/broadcast.jl")
 
+    include_test("tuplecast.jl")  # used primarily for broadcast
+
     println()
 
     include_test("rulesets/Statistics/statistics.jl")

test/test_helpers.jl

@@ -98,6 +98,7 @@ function _gpu_test(::typeof(frule), f::Function, g::Function, xs...; kw...)  # s
     _gpu_test(frule, xdots, f, g, xs...; kw...)
 end
 
+const CFG = ChainRulesTestUtils.TestConfig()
 
 """
     Multiplier(x)
@@ -181,8 +182,8 @@ fstar(A, B) = A * B
 ChainRulesCore.frule((_, ΔA, ΔB), ::typeof(fstar), A, B) = A * B, muladd(ΔA, B, A * ΔB)
 
 "A version of `log` with only an `frule` defined"
-flog(x:::Number) = log(x)
-ChainRulesCore.frule((_, xdot), ::typeof(flog), x::Number) = log(x), inv(x) * xdot
+flog(x::Number) = log(x)
+ChainRulesCore.frule((_, Δx), ::typeof(flog), x::Number) = log(x), inv(x) * Δx
 
 @testset "test_helpers.jl" begin