Conditionals do not work in array comprehensions

[gregoryquick]

Can get error with:

A = [0.0 1.0 2.0]
f(x) = sum([x[i] for i in 1:3 if i != 100]) 
Zygote.gradient(f, A) 

Which gives:

MethodError: Cannot `convert` an object of type Nothing to an object of type Int64

Backtrace

    setindex!(::Vector{Int64}, ::Nothing, ::Int64)@array.jl:903
    [email protected]:317[inlined]
    copyto_unaliased!(::IndexCartesian, ::SubArray{Int64, 1, Vector{Int64}, Tuple{Vector{Int64}}, false}, ::IndexLinear, ::Vector{Nothing})@abstractarray.jl:1032
    [email protected]:998[inlined]
    [email protected]:954[inlined]
    [email protected]:913[inlined]
    [email protected]:871[inlined]
    [email protected]:868[inlined]
    (::Zygote.var"#606#607"{Vector{Int64}, Vector{Bool}})(::Vector{Nothing})@array.jl:264
    #2661#[email protected]:67[inlined]
    #[email protected]:41[inlined]
    (::Zygote.var"#2685#back#613"{Zygote.var"#57#58"{Zygote.var"#2661#back#608"{Zygote.var"#606#607"{Vector{Int64}, Vector{Bool}}}}})(::Vector{Nothing})@adjoint.jl:67
    Pullback@Other: 1[inlined]
    (::typeof(∂(f)))(::Float64)@interface2.jl:0
    (::Zygote.var"#57#58"{typeof(∂(f))})(::Float64)@interface.jl:41
    gradient(::Function, ::Matrix{Float64})@interface.jl:76
    top-level scope@Local: 1[inlined]

[mcabbott]

Looks like the bug is in the gradient of filter, array.jl:264, which after #785 is what Iterators.filter (which is what the conditional comprehension really is) calls to do the work. Another way to trigger it is:

julia> gradient(x -> sum(map(i -> x[i], filter(i -> i != 100, 1:3))), A)
ERROR: MethodError: Cannot `convert` an object of type Nothing to an object of type Int64

[mtfishman]

I also see:

julia> x = 3
3

julia> f1(x) = sum(Int[x^2 for i in 1:2])
f1 (generic function with 1 method)

julia> f1'(x)
ERROR: Mutating arrays is not supported -- called setindex!(::Vector{Int64}, _...)
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:33
 [2] (::Zygote.var"#442#443"{Vector{Int64}})(#unused#::Nothing)
   @ Zygote ~/.julia/packages/Zygote/FPUm3/src/lib/array.jl:71
 [3] (::Zygote.var"#2339#back#444"{Zygote.var"#442#443"{Vector{Int64}}})(Δ::Nothing)
   @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
 [4] Pullback
   @ ./REPL[102]:1 [inlined]
 [5] (::typeof(∂(f1)))(Δ::Int64)
   @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [6] #57
   @ ~/.julia/packages/Zygote/FPUm3/src/compiler/interface.jl:41 [inlined]
 [7] (::Zygote.var"#59#60"{typeof(f1)})(x::Int64)
   @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface.jl:83
 [8] top-level scope
   @ REPL[103]:1

and:

julia> f2(x) = sum([x^2 for i in 1:2 for j in 1:2])
f2 (generic function with 1 method)

julia> f2'(x)
ERROR: Mutating arrays is not supported -- called push!(::Vector{Int64}, _...)
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:33
  [2] (::Zygote.var"#450#451"{Vector{Int64}})(#unused#::Nothing)
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/lib/array.jl:78
  [3] (::Zygote.var"#2359#back#452"{Zygote.var"#450#451"{Vector{Int64}}})(Δ::Nothing)
    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
  [4] Pullback
    @ ./array.jl:823 [inlined]
  [5] (::typeof(∂(grow_to!)))(Δ::FillArrays.Fill{Int64, 1, Tuple{Base.OneTo{Int64}}})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
  [6] Pullback
    @ ./array.jl:801 [inlined]
  [7] (::typeof(∂(grow_to!)))(Δ::FillArrays.Fill{Int64, 1, Tuple{Base.OneTo{Int64}}})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
  [8] Pullback
    @ ./array.jl:701 [inlined]
  [9] (::typeof(∂(_collect)))(Δ::FillArrays.Fill{Int64, 1, Tuple{Base.OneTo{Int64}}})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [10] Pullback
    @ ./array.jl:649 [inlined]
 [11] (::typeof(∂(collect)))(Δ::FillArrays.Fill{Int64, 1, Tuple{Base.OneTo{Int64}}})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [12] Pullback
    @ ./REPL[104]:1 [inlined]
 [13] (::typeof(∂(f2)))(Δ::Int64)
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [14] (::Zygote.var"#57#58"{typeof(∂(f2))})(Δ::Int64)
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface.jl:41
 [15] (::Zygote.var"#59#60"{typeof(f2)})(x::Int64)
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface.jl:83
 [16] top-level scope
    @ REPL[105]:1

[mcabbott]

These are I think independent problems with comprehensions. Typed comprehensions go their own path and don't seem to work at all:

julia> gradient(xs -> sum([x^2 for x in xs]), [1,2,3])
([2.0, 4.0, 6.0],)

julia> gradient(xs -> sum(Float32[x^2 for x in xs]), [1,2,3])
ERROR: Mutating arrays is not supported -- called setindex!(::Vector{Float32}, _...)

And #785 made no attempt to handle Iterators.flatten, which is what backs the for for pattern:

julia> gradient(xs -> sum([i/j for i in xs for j in xs]), [1,2,3])
ERROR: Mutating arrays is not supported -- called push!(::Vector{Float64}, _...)

Maybe I'd call both missing rules, not bugs in existing rules.

[ToucheSir]

Typed comprehensions lower to getindex(::Type, ::Generator), I believe? Never seen that come up before so I'm not surprised by the error.