SciML / SciML/ComponentArrays.jl

Construction of ComponentArray inside of AD/Zygote

Open
#176 2 comments 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.