reducedim_init MethodError with structs

[cossio]

using Zygote
struct A
    x::Float64
end
f(a,v) = a.x + v
g(A,V) = sum(f.(A,V))
gradient(g, A.(randn(2)), randn(2,2)) # Error

Error message and stacktrace:

ERROR: LoadError: MethodError: no method matching reducedim_init(::typeof(identity), ::typeof(Zygote.accum), ::Array{NamedTuple{(:x,),Tuple{Float64}},2}, ::Tuple{Int64,Int64})
Closest candidates are:
  reducedim_init(::Any, ::Union{typeof(+), typeof(Base.add_sum)}, ::AbstractArray, ::Any) at reducedim.jl:109
  reducedim_init(::Any, ::Union{typeof(*), typeof(Base.mul_prod)}, ::AbstractArray, ::Any) at reducedim.jl:112
  reducedim_init(::Any, ::typeof(min), ::AbstractArray, ::Any) at reducedim.jl:132
  ...
Stacktrace:
 [1] _mapreduce_dim(::Function, ::Function, ::NamedTuple{(),Tuple{}}, ::Array{NamedTuple{(:x,),Tuple{Float64}},2}, ::Tuple{Int64,Int64}) at ./reducedim.jl:317
 [2] #mapreduce#580 at ./reducedim.jl:307 [inlined]
 [3] #reduce#582 at ./reducedim.jl:352 [inlined]
 [4] accum_sum(::Array{NamedTuple{(:x,),Tuple{Float64}},2}; dims::Tuple{Int64,Int64}) at /home/cossio/.julia/packages/Zygote/4tJp5/src/lib/broadcast.jl:36
 [5] unbroadcast at /home/cossio/.julia/packages/Zygote/4tJp5/src/lib/broadcast.jl:48 [inlined]
 [6] map(::typeof(Zygote.unbroadcast), ::Tuple{Array{Main.ZygoteStruct.A,1},Array{Float64,2}}, ::Tuple{Array{NamedTuple{(:x,),Tuple{Float64}},2},Array{Float64,2}}) at ./tuple.jl:177
 [7] (::Zygote.var"#1734#1741"{Tuple{Array{Main.ZygoteStruct.A,1},Array{Float64,2}},Val{3},Array{typeof(∂(f)),2}})(::FillArrays.Fill{Float64,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}) at /home/cossio/.julia/packages/Zygote/4tJp5/src/lib/broadcast.jl:139
 [8] #4367#back at /home/cossio/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49 [inlined]
 [9] (::Zygote.var"#175#176"{Zygote.var"#4367#back#1745"{Zygote.var"#1734#1741"{Tuple{Array{Main.ZygoteStruct.A,1},Array{Float64,2}},Val{3},Array{typeof(∂(f)),2}}},Tuple{NTuple{4,Nothing},Tuple{}}})(::FillArrays.Fill{Float64,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}) at /home/cossio/.julia/packages/Zygote/4tJp5/src/lib/lib.jl:170
 [10] #344#back at /home/cossio/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49 [inlined]
 [11] broadcasted at ./broadcast.jl:1238 [inlined]
 [12] g at /home/cossio/Desktop/julia/zygote_structarray (copy).jl:8 [inlined]
 [13] (::typeof(∂(g)))(::Float64) at /home/cossio/.julia/packages/Zygote/4tJp5/src/compiler/interface2.jl:0
 [14] (::Zygote.var"#36#37"{typeof(∂(g))})(::Float64) at /home/cossio/.julia/packages/Zygote/4tJp5/src/compiler/interface.jl:36
 [15] gradient(::Function, ::Array{Main.ZygoteStruct.A,1}, ::Vararg{Any,N} where N) at /home/cossio/.julia/packages/Zygote/4tJp5/src/compiler/interface.jl:45
 [16] top-level scope at /home/cossio/Desktop/julia/zygote_structarray (copy).jl:10
 [17] include(::String) at ./client.jl:439
 [18] top-level scope at REPL[1]:1
in expression starting at /home/cossio/Desktop/julia/zygote_structarray (copy).jl:10

I originally reported the error using StructArray (see JuliaArrays/StructArrays.jl#124). But later realized the same issue occurs with ordinary structs so I simplified the example.

[cossio]

It looks like similar errors have been walked-around by defining an appropriate accum_sum, as in

Zygote.jl/src/lib/broadcast.jl

Line 43 in 855a158

accum_sum(xs::AbstractArray{<:AbstractArray{<:Number}}; dims = :) = sum(xs, dims = dims)

.

[CarloLucibello]

Did you try to define accum_sum(::Array{NamedTuple}; dims)? Also, I think that accum_sum should receive a StructArray here, something in map or unbroadcast should be changed

[cossio]

Thanks for taking a look.

Did you try to define accum_sum(::Array{NamedTuple}; dims)?

I'm not sure how to.

Also, I think that accum_sum should receive a StructArray here, something in map or unbroadcast should be changed

Yes I agree. It seems inefficient to convert the StructArray to tuples, and then reassemble back.

[CarloLucibello]

Could you try if something like this helps?

julia> function _add(x1::T,x2::T) where T<:NamedTuple{names,types} where {names, types}
         vals = ((getfield(x1,f) + getfield(x2,f) for f in names)...,)
         NamedTuple{names,types}(vals)
       end
_add (generic function with 1 method)

julia> accum_sum(xs::AbstractArray{<:NamedTuple}; dims = :) = reduce(_add, xs; dims=dims)
accum_sum (generic function with 3 methods)

julia> x = (a=1,b=2.)
(a = 1, b = 2.0)

julia> accum_sum([x,x,x])
(a = 3, b = 6.0)

[cossio]

Could you try if something like this helps?

I get the same MethodError: no method matching reducedim_init error after defining Zygote.accum_sum this way.

[cossio]

Defining

function Base.reducedim_init(f::typeof(identity), op::typeof(Zygote.accum), A::AbstractArray, region)
    return Base.reducedim_initarray(A, region, nothing, Union{Nothing,eltype(A)})
end

works. Thanks @MikeInnes for guidance.

If there are some improvements you can think of let me know. Should I make a PR with this?

[CarloLucibello]

It's not typestable though, are we still happy with it?

julia> gradient(g, A.(randn(2)), randn(2,2))[1] 
2-element reshape(::Array{Union{Nothing, NamedTuple{(:x,),Tuple{Float64}}},2}, 2) with eltype Union{Nothing, NamedTuple{(:x,),Tuple{Float64}}}:
 (x = 2.0,)
 (x = 2.0,)

[cossio]

It's not typestable though, are we still happy with it?

Because of the Union{Nothing,...}? The problem is that the generic zero is nothing in Zygote, which has to be combined with eltype(A) that can be anything. I didn't know of a better way to go around this.

[DhairyaLGandhi]

I think we should generally not try to override methods that can change within minor releases in Julia

[MikeInnes]

A small union is not that bad performance-wise; in any case it's a strict improvement over throwing an error. It'd be good to figure out a type-stable approach once this fix is in.