# Accessing a ComponentArray's fields in a type-stable manner

**URL:** <https://discourse.julialang.org/t/accessing-a-componentarrays-fields-in-a-type-stable-manner/103832>\
**Category:** General Usage\
**Tags:** namedtuple, componentarrays, lux\
**Created:** [September 13, 2023, 9:18pm UTC](https://discourse.julialang.org/t/accessing-a-componentarrays-fields-in-a-type-stable-manner/103832 "2023-09-13T21:18:12Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)\
**Post date:** [September 13, 2023, 9:18pm UTC](https://discourse.julialang.org/t/accessing-a-componentarrays-fields-in-a-type-stable-manner/103832/1 "2023-09-13T21:18:12Z")

</div>

I’m struggling with iterating over the named fields of a ComponentArray in a compiler-friendly way.

Take the following MWE, emulating some Neural Network:

1. I create the nested structure `layers`, containing some weight matrices. Each is itself a `NamedTuple` with a lower-triangular and an upper-triangular component.

2. I feed it to the function `feedforward`, which applies each layer to the input in sequence.

For `layers` a `NamedTuple`, I can simply iterate over the elements of `layers` and it works relatively well:

> **Code iterating over a NamedTuple**
>
> ```julia
> using NamedTupleTools, LinearAlgebra, ComponentArrays
> 
> function create_layers(size,number) 
> l = (lower_weight = LowerTriangular(rand(size,size)),
> upper_weight = UpperTriangular(rand(size,size)))
> 
> layer_keys = Symbol.(["layer_$i" for i in 1:number])
> layer_data = [l for _ in 1:number]
> layers = namedtuple(layer_keys,layer_data)
> return layers
> end 
> 
> function feedforward(x::AbstractArray, layers)
> fwd = x
> for L in layers
> fwd = (L.lower_weight*L.upper_weight*fwd)
> end     
> return fwd 
> end 
> 
> layers = create_layers(3,3)
> input = rand(3)
> feedforward(input,layers) #Works OK
> 
> ```

However, when I turn `layers` into a `ComponentArray`, indexing into it now returns its individual elements, breaking everything:

> **Incorrect attempt with ComponentArrays**
>
> ```julia
> layers_CA = ComponentArray(layers)
> feedforward(input,layers_CA) #ERROR: type Float64 has no field lower_weight
> 
> ```

A possible workaround is to collect `layers_CA`’s keys and then iterate over them, but then the compiler can’t really prepare for what’s coming and the result is type instability:

> **Code iterating over keys**
>
> ```julia
> function feedforward(x::AbstractArray, ps::ComponentArray)
> fwd = x
> for k in keys(ps)
> L = ps[k]
> fwd = (L.lower_weight*L.upper_weight*fwd)
> end     
> return fwd 
> end
> 
> feedforward(input,layers_CA) #Works
> @code_warntype feedforward(input,layers_CA)
> 
> ```

Resulting in return-type of Any:

> **Result of @code\_warntype**
>
> ```julia
> MethodInstance for feedforward(::Vector{Float64}, ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:18, Axis(lower_weight = ViewAxis(1:9, ShapedAxis((3, 3), NamedTuple())), upper_weight = ViewAxis(10:18, ShapedAxis((3, 3), NamedTuple())))), layer_2 = ViewAxis(19:36, Axis(lower_weight = ViewAxis(1:9, ShapedAxis((3, 3), NamedTuple())), upper_weight = ViewAxis(10:18, ShapedAxis((3, 3), NamedTuple())))), layer_3 = ViewAxis(37:54, Axis(lower_weight = ViewAxis(1:9, ShapedAxis((3, 3), NamedTuple())), upper_weight = ViewAxis(10:18, ShapedAxis((3, 3), NamedTuple())))))}}})
> from feedforward(x::AbstractArray, ps::ComponentArray) @ Main Untitled-4:44
> Arguments
> #self#::Core.Const(feedforward)
> x::Vector{Float64}
> ps::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:18, Axis(lower_weight = ViewAxis(1:9, ShapedAxis((3, 3), NamedTuple())), upper_weight = ViewAxis(10:18, ShapedAxis((3, 3), NamedTuple())))), layer_2 = ViewAxis(19:36, Axis(lower_weight = ViewAxis(1:9, ShapedAxis((3, 3), NamedTuple())), upper_weight = ViewAxis(10:18, ShapedAxis((3, 3), NamedTuple())))), layer_3 = ViewAxis(37:54, Axis(lower_weight = ViewAxis(1:9, ShapedAxis((3, 3), NamedTuple())), upper_weight = ViewAxis(10:18, ShapedAxis((3, 3), NamedTuple())))))}}}
> Locals
> @_4::Union{Nothing, Tuple{Symbol, Int64}}
> fwd::Any
> k::Symbol
> L::Any
> Body::Any
> 1 ─ (fwd = x)
> │ %2 = Main.keys(ps)::Core.Const((:layer_1, :layer_2, :layer_3))
> │ (@_4 = Base.iterate(%2))
> │ %4 = (@_4::Core.Const((:layer_1, 2)) === nothing)::Core.Const(false)
> │ %5 = Base.not_int(%4)::Core.Const(true)
> └── goto #4 if not %5
> 2 ┄ %7 = @_4::Tuple{Symbol, Int64}
> │ (k = Core.getfield(%7, 1))
> │ %9 = Core.getfield(%7, 2)::Int64
> │ (L = Base.getindex(ps, k))
> │ %11 = Base.getproperty(L, :lower_weight)::Any
> │ %12 = Base.getproperty(L, :upper_weight)::Any
> │ (fwd = %11 * %12 * fwd)
> │ (@_4 = Base.iterate(%2, %9))
> │ %15 = (@_4 === nothing)::Bool
> │ %16 = Base.not_int(%15)::Bool
> └── goto #4 if not %16
> 3 ─ goto #2
> 4 ┄ return fwd
> 
> ```

How should I go about this? How can I create a function that knows (at compile time) what keys I’m going to iterate over without coding the function myself for each different layout of `layers`? I’ve seen people do this with metaprogramming but I don’t feel experienced enough to use `@generated` functions and the like.

I need this for a few custom `Lux` layers I’m constructing, which need to be easily translatable into `ComponentArrays` to be compatible with NeuralPDE.

---

<div class="post-metadata">

**Author:** ![jonniedie](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jonniedie/32/12842_2.png) [@jonniedie](https://discourse.julialang.org/u/jonniedie)\
**Post date:** [October 31, 2023, 3:41am UTC](https://discourse.julialang.org/t/accessing-a-componentarrays-fields-in-a-type-stable-manner/103832/2 "2023-10-31T03:41:33Z")

</div>

ComponentArrays exports a function called `valkeys` that’s like `keys` but with each key wrapped in a `Val` for this exact reason. Try

```julia
function feedforward(x::AbstractArray, ps::ComponentArray)
    fwd = x
    for k in valkeys(ps)
        L = @view ps[k]
        fwd = (L.lower_weight*L.upper_weight*fwd)
    end     
    return fwd 
end

```

---

<div class="post-metadata">

**Author:** ![jonniedie](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jonniedie/32/12842_2.png) [@jonniedie](https://discourse.julialang.org/u/jonniedie)\
**Post date:** [October 31, 2023, 3:53am UTC](https://discourse.julialang.org/t/accessing-a-componentarrays-fields-in-a-type-stable-manner/103832/3 "2023-10-31T03:53:15Z")

</div>

Another option is to do `for L in NamedTuple(ps)` for the iteration. This seems to make `@code_warntype` slightly happier than the `valkeys` solution, but doesn’t make the `@btime` performance (time or allocations) any better or worse.