How to get index range of subarray?

[vpuri3]

In the case below, how do I get the range of indices spanned by x.a?

x = (; a = ones(2), b = ones(2)) |> ComponentArray

idx_a = ?? # want 1:2

[vpuri3]

@jonniedie is that possible? TIA

[jonniedie]

It's a little clunky, but you can do

only(getaxes(x))[:a].idx

[vpuri3]

Thanks, that works. Is there a similar way to get the global indices of nested arrays? eg.

ca = (; a = ones(2), b = (; x = 2 * ones(2), y = 3 * ones(2))) |> ComponentArray
idx_b_x = getindices(ca.b.x) # want 3:4

[alicedb2]

This would be extremely useful for my current use-case involving neural networks with Lux and DiffEqFlux. Using ComponentArrays is very useful for manipulating the parameter vector of a whole NN which for k dense layers you'd have something like

p = (layer_1=(weight=[1 3; 2 4], bias=[5, 6]), layer_2=(weight=[7 9; 8 10], bias=[11, 12]), ..., layer_k=...)

Right now I would need to iterate over the layers in an arbitrary nn and change specific columns of the weight matrix of some given layer which satisfy some condition at runtime and would need to specify using something which is not an explicit p.layer_2.weight[:, 1] .= 0, or if I create some matrix of the form p * p' I'd like to be able to index a nested submatrix and do some operations on it.

Getting the global indices of an arbitrary nested subarray would solve this.