Fix adjoint Iterators.product behavior with nothing (#1170)

* Fix adjoint Iterators.product behavior with nothing

* Apply suggestions from code review

Co-authored-by: Michael Abbott <[email protected]>

* add Iterators.product adjoint tests

* Update test/lib/array.jl

Co-authored-by: Michael Abbott <[email protected]>

Co-authored-by: Michael Abbott <[email protected]>

Author: lxvm

src/lib/array.jl

@@ -286,10 +286,10 @@ _ndims(x) = Base.IteratorSize(x) isa Base.HasShape ? _ndims(Base.IteratorSize(x)
   function back(dy::AbstractArray)
     d = 1
     ntuple(length(xs)) do n
-      first(dy)[n] === nothing && return nothing
       nd = _ndims(xs[n])
       dims = ntuple(i -> i<d ? i : i+nd, ndims(dy)-nd)
       d += nd
+      first(dy)[n] === nothing && return nothing
       init = zero.(first(dy)[n]) # allows for tuples, which accum can add:
       red = mapreduce(StaticGetter{n}(), accum, dy; dims=dims, init=init)
       return _project(xs[n], reshape(red, axes(xs[n])))

test/lib/array.jl

@@ -1,8 +1,20 @@
 using ChainRulesTestUtils
 using LinearAlgebra
-using Zygote: ZygoteRuleConfig
+using Zygote: ZygoteRuleConfig, _pullback
 
 # issue 897
 
 test_rrule(ZygoteRuleConfig(), x->sum(sin, Diagonal(x)), ones(2); rrule_f=rrule_via_ad, check_inferred=false)
 test_rrule(ZygoteRuleConfig(), x->sum(sin, Diagonal(x)), rand(3); rrule_f=rrule_via_ad, check_inferred=false)
+
+@testset "adjoints of Iterators.product" begin
+    y, back = _pullback(Iterators.product, 1:5, 1:3, 1:2)
+    @test back(collect(y)) == (nothing, [6.0, 12.0, 18.0, 24.0, 30.0], [10.0, 20.0, 30.0], [15.0, 30.0])
+    @test back([(nothing, j, k) for i in 1:5, j in 1:3, k in 1:2]) == (nothing, nothing, [10.0, 20.0, 30.0], [15.0, 30.0])
+    @test back([(i, nothing, k) for i in 1:5, j in 1:3, k in 1:2]) == (nothing, [6.0, 12.0, 18.0, 24.0, 30.0], nothing, [15.0, 30.0])
+    @test back([(i, j, nothing) for i in 1:5, j in 1:3, k in 1:2]) == (nothing, [6.0, 12.0, 18.0, 24.0, 30.0], [10.0, 20.0, 30.0], nothing)
+
+    # This was wrong before https://github.com/FluxML/Zygote.jl/pull/1170
+    @test gradient(x -> sum([y[2] * y[3] for y in Iterators.product(x, x, x, x)]), [1,2,3,4])[1] ≈ [320, 320, 320, 320]
+    @test gradient(x -> sum(y[2] * y[3] for y in Iterators.product(x, x, x, x)), [1,2,3,4])[1] ≈ [320, 320, 320, 320]
+end