Merge pull request #1525 from BioTurboNick/enumerate-cond-gen

Support `enumerate` in conditional generators/comprehensions

src/lib/array.jl

@@ -260,6 +260,7 @@ end
 @adjoint function enumerate(xs)
   back(::AbstractArray{Nothing}) = nothing
   back(dy::NamedTuple{(:itr,)}) = tuple(dy.itr)
+  back(diys::AbstractArray{Union{Nothing, T}}) where T = (map(x -> x === nothing ? x : last(x), diys),)
   back(diys) = (map(last, diys),)
   enumerate(xs), back
 end

test/lib/array.jl

@@ -118,6 +118,19 @@ end
         sum(v for (_, v) in d)
     end
     @test gradient(f_comprehension, w)[1] == ones(5)
+
+    w = [randn(5); NaN]
+    function f_generator_conditional(w)
+        d = Dict{Int, Float64}(i => v for (i,v) in enumerate(w) if !isnan(v))
+        sum(v for (_, v) in d)
+    end
+    @test gradient(f_generator_conditional, w)[1] == [ones(5); nothing]
+
+    function f_comprehension_conditional(w)
+        d = Dict{Int, Float64}(i => v for (i,v) in enumerate(w) if !isnan(v))
+        sum(v for (_, v) in d)
+    end
+    @test gradient(f_comprehension_conditional, w)[1] == [ones(5); nothing]
 end
 
 @testset "_reverse" begin