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