SciML / SciML/ComponentArrays.jl
Construction of ComponentArray inside of AD/Zygote
Nobody has claimed this yet.
- Dominant language
- Julia
- Stars
- 375
- Forks
- 42
- Avg merge
- 7h 25m
- Merged PRs (30d)
- 17
Description
I want to compute the gradient of a loss function with respect to a ComponentArray. In the loss function, I need to reconstruct a ComponentArray.
Based on @jonniedie reply https://github.com/jonniedie/ComponentArrays.jl/issues/126#issuecomment-1141580528, I tried
function my_sum(v)
ax = getaxes(v)
@unpack x, y = v
ca = ComponentArray([x..., y...], ax)
return sum(ca.x + ca.y)
end
Zygote.gradient(my_sum, ComponentArray(x=[0.0], y=[0.0]))
which fails with
ERROR: ArgumentError: indexed assignment with a single value to possibly many locations is not supported; perhaps use broadcasting `.=` instead?
Stacktrace:
[1] setindex_shape_check(::ChainRulesCore.Tangent{Any, Tuple{Float64}}, ::Int64)
@ Base ./indices.jl:261
[2] _unsafe_setindex!(#unused#::IndexLinear, A::Vector{Float64}, x::ChainRulesCore.Tangent{Any, Tuple{Float64}}, I::UnitRange{Int64})
@ Base ./multidimensional.jl:939
[3] _setindex!
@ ./multidimensional.jl:930 [inlined]
[4] setindex!
@ ./abstractarray.jl:1344 [inlined]
[5] macro expansion
@ ~/.julia/packages/ComponentArrays/EjZNJ/src/array_interface.jl:0 [inlined]
[6] _setindex!(x::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(x = 1:1, y = 2:2)}}}, v::ChainRulesCore.Tangent{Any, Tuple{Float64}}, idx::Val{:y})
@ ComponentArrays ~/.julia/packages/ComponentArrays/EjZNJ/src/array_interface.jl:129
[7] setproperty!
@ ~/.julia/packages/ComponentArrays/EjZNJ/src/namedtuple_interface.jl:17 [inlined]
[8] (::ComponentArrays.var"#getproperty_adjoint#87"{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(x = 1:1, y = 2:2)}}}, Symbol})(Δ::ChainRulesCore.Tangent{Any, Tuple{Float64}})
@ ComponentArrays ~/.julia/packages/ComponentArrays/EjZNJ/src/compat/chainrulescore.jl:4
[9] ZBack
@ ~/.julia/packages/Zygote/PD12J/src/compiler/chainrules.jl:206 [inlined]
[10] Pullback
@ ~/.julia/packages/UnPack/EkESO/src/UnPack.jl:34 [inlined]
[11] (::typeof(∂(unpack)))(Δ::Tuple{Float64})
@ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
[12] macro expansion
@ ~/.julia/packages/UnPack/EkESO/src/UnPack.jl:101 [inlined]
[13] Pullback
pointing to the @unpack call. @avik-pal noted that it also happens even without the @unpack
function my_sum(v)
ax = getaxes(v)
ca = ComponentArray([v.x..., v.y...], ax)
return sum(ca.x + ca.y)
end
Zygote.gradient(my_sum, ComponentArray(x=[0.0], y=[0.0]))
but is resolved by using vcat
function my_sum(v)
ax = getaxes(v)
@unpack x, y = v
ca = ComponentArray(vcat(x,y), ax)
return sum(ca.x + ca.y)
end
The issue seems to be that \Delta is a Tuple{Float64} in
https://github.com/jonniedie/ComponentArrays.jl/blob/cbb24ef7156d18f1576ea48d7ae42023cc5bfa70/src/compat/chainrulescore.jl#L4
for splatting.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with src/compat/chainrulescore.jl at the referenced line and reproduce the failure using the two my_sum examples with Zygote.gradient. Compare the splatted array construction with the working vcat form and inspect how the Tangent delta is handled. Done means reconstructing a ComponentArray inside the gradient path works without the indexed-assignment error, with a regression test for the example.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100