AD failing with ComponentArrays

[gdalle]

Hi there, and thanks for the great package
Are ComponentArrays supposed to be AD-compatible? If so, did I do something wrong with the following code?

julia> using ComponentArrays, Zygote

julia> function mysum(x::AbstractVector)
           y = ComponentVector(x=x)
           return sum(y)
       end
mysum (generic function with 1 method)

julia> Zygote.gradient(mysum, rand(10))
ERROR: MethodError: no method matching size(::NamedTuple{(:x,), Tuple{Vector{Float64}}})

[jonniedie]

Hey, sorry for the delay. I really hope to get to this tonight.

[gdalle]

No worries, it's not urgent! Just tell me if there is anything I can do to help

[roflmaostc]

Hi!

I stumbled over the same issue today. Is there any progress?

Best,

Felix

[jonniedie]

Hey. Sorry to get back so late to this. I've been swamped with work and honestly just haven't had time to get to it. I doubt I'll be able to get to it before this weekend. If anyone has some experience with AD through constructors in Zygote, I'd appreciate the help. If not, I'm shooting for the weekend to look at it.

[roflmaostc]

Thanks a lot!
But no hurry, I just was interested in the progress 😆

Because of #48 it might be also a bad idea to put the creation in the forward model.

[jonniedie]

So I gave this a shot and it ultimately isn't going to work with the way ComponentArrays are currently built from keyword or NamedTuples due to Zygote's limitation on mutation. It would probably be possible to make this version of construction not use mutation, but I'm not sure that will be too worthwhile for practical uses, since this style of construction is already a bit too slow to be useful in most cases due to its inherent type instability.

[gdalle]

Thanks for taking a look! Which construction mode would you suggest for cases where efficiency is paramount?

…

Le 28 mai 2022 à 21:02, Jonnie Diegelman ***@***.***> a écrit :  So I gave this a shot and it ultimately isn't going to work with the way ComponentArrays are currently built from keyword or NamedTuples due to Zygote's limitation on mutation. It would probably be possible to make this version of construction not use mutation, but I'm not sure that will be too worthwhile for practical uses, since this style of construction is already a bit too slow to be useful in most cases due to its inherent type instability. — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you authored the thread.

[jonniedie]

I'll usually use the ComponentArray(data::AbstractArray, axis::AbstractAxis) style of construction. So like

ca_prototype = ComponentArray(x=0.0, y=0.0)
const ax = getaxes(ca_prototype)

function my_sum(v)
    ca = ComponentArray(v, ax)
    return ca.x + ca.y
end

Zygote.gradient(my_sum, rand(2))

[gdalle]

Great, thanks! I think it would be nice to add a line or two to the docs describing this constructor: it is briefly mentioned in the API but never demonstrated. Maybe also add a warning related to AD with Zygote?

[jonniedie]

That's probably a good idea. The docs could also use a bit of work in general.