# Indexing an arbitrary array over its last dimension

**URL:** <https://discourse.julialang.org/t/indexing-an-arbitrary-array-over-its-last-dimension/70443>\
**Category:** General Usage\
**Tags:** array, zygote\
**Created:** [October 27, 2021, 12:20am UTC](https://discourse.julialang.org/t/indexing-an-arbitrary-array-over-its-last-dimension/70443 "2021-10-27T00:20:11Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![MattSainsbury-Dale](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mattsainsbury-dale/32/221714_2.png) [@MattSainsbury-Dale](https://discourse.julialang.org/u/MattSainsbury-Dale)\
**Post date:** [October 27, 2021, 12:20am UTC](https://discourse.julialang.org/t/indexing-an-arbitrary-array-over-its-last-dimension/70443/1 "2021-10-27T00:20:11Z")

</div>

I am trying to index over the last dimension of an array of unknown dimension. Please see the following MWE for a 2-dimensional array:

```julia
a = rand(2, 4)
indices = [1:2, 3:4]
colons = repeat([:], ndims(a) - 1)
v = [a[colons..., idx] for idx ∈ indices]

```

This works as intended:

```julia
# Sanity check: 
v[1] == a[:, 1:2] # true
v[2] == a[:, 3:4] # true

```

However, I received an error when using this approach while training a `Flux` model:

```julia
ERROR: MethodError: no method matching reshape(::Tuple{Nothing}, ::Int64, ::Colon)
Closest candidates are:
  reshape(::FillArrays.AbstractFill, ::Union{Colon, Int64}...) at /home/matthew/.julia/packages/FillArrays/Vzxer/src/FillArrays.jl:221
  reshape(::FillArrays.AbstractFill, ::Union{Colon, Integer}...) at /home/matthew/.julia/packages/FillArrays/Vzxer/src/FillArrays.jl:222
  reshape(::ChainRulesCore.AbstractZero, ::Any...) at /home/matthew/.julia/packages/ChainRulesCore/7OROc/src/tangent_types/abstract_zero.jl:37
  ...
Stacktrace:
  [1] (::Zygote.var"#542#543"{Vector{Colon}})(ȳ::Tuple{Nothing})
    @ Zygote ~/.julia/packages/Zygote/l3aNG/src/lib/array.jl:168
  [2] (::Zygote.var"#2555#back#544"{Zygote.var"#542#543"{Vector{Colon}}})(Δ::Tuple{Nothing})
    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
  [3] Pullback
    @ ~/Dropbox/SpatialDeepSets/src/DeepSetStruct.jl:163 [inlined]
  [4] (::typeof(∂(λ)))(Δ::CuArray{Float32, 2})
    @ Zygote ~/.julia/packages/Zygote/l3aNG/src/compiler/interface2.jl:0
  [5] Pullback
    @ ~/Dropbox/SpatialDeepSets/src/Train.jl:116 [inlined]
  [6] (::typeof(∂(λ)))(Δ::Float32)
    @ Zygote ~/.julia/packages/Zygote/l3aNG/src/compiler/interface2.jl:0
  [7] (::Zygote.var"#90#91"{Params, typeof(∂(λ)), Zygote.Context})(Δ::Float32)
    @ Zygote ~/.julia/packages/Zygote/l3aNG/src/compiler/interface.jl:348
  [8] gradient(f::Function, args::Params)
    @ Zygote ~/.julia/packages/Zygote/l3aNG/src/compiler/interface.jl:76
  [9] train(model::DeepSet{Chain{Tuple{Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, typeof(Flux.flatten), Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, typeof(mean_over_last_dim), typeof(identity)}, data_simulator::typeof(data_simulator), train_params::ParameterConfigurations, val_params::ParameterConfigurations; run::Int64, save_best_model_only::Bool, args::Args, num_prior_epochs::Int64)
    @ Main ~/Dropbox/SpatialDeepSets/src/Train.jl:116
 [10] top-level scope
    @ ~/Dropbox/SpatialDeepSets/src/GaussianProcess/nuVaried/TrainNF30.jl:31
 [11] eval
    @ ./boot.jl:360 [inlined]

```

By hard coding the number of dimensions for a specific case, I have confirmed that it is my above indexing that is causing this error; that is, my code works if I use `a[:, idx]` rather than `a[colons..., idx]`.

Is there a better way to subset an array with unknown `ndims` over its last dimension?

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [October 27, 2021, 1:00am UTC](https://discourse.julialang.org/t/indexing-an-arbitrary-array-over-its-last-dimension/70443/2 "2021-10-27T01:00:45Z")

</div>

It’s a Zygote bug, that ought to work. Please make an issue:

```julia
julia> gradient(a -> sum(a[repeat([:], ndims(a) - 1)..., 1]), rand(3,3,3))
ERROR: MethodError: no method matching reshape(::Tuple{Nothing, Nothing}, ::Int64, ::Colon)
Closest candidates are:
  reshape(::AbstractZero, ::Any...) at ~/.julia/packages/ChainRulesCore/Z8WET/src/tangent_types/abst

```

However, you should probably use `selectdim` here. Or write `ntuple(_ -> (:), ndims(a) - 1)` here, making a tuple not an array.

---

<div class="post-metadata">

**Author:** ![MattSainsbury-Dale](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mattsainsbury-dale/32/221714_2.png) [@MattSainsbury-Dale](https://discourse.julialang.org/u/MattSainsbury-Dale)\
**Post date:** [October 27, 2021, 2:37am UTC](https://discourse.julialang.org/t/indexing-an-arbitrary-array-over-its-last-dimension/70443/3 "2021-10-27T02:37:51Z")

</div>

Thank you, `selectdim()` is the neat function that I was looking for. However, it also caused problems during training because it seemingly does some scalar indexing under the hood (I’m working with the GPU, which doesn’t like scalar indexing). Fortunately, the `ntuple()` approach works without issue.

I was very surprised that splatting the colons from a `tuple` gives different results to splatting them from an `array`. I had thought that splatting involved throwing all elements of a container into the function, and so the container type doesn’t matter.

I’ve opened an issue with `Zygote` [here](https://github.com/FluxML/Zygote.jl/issues/1110).

---

<div class="post-metadata">

**Author:** ![tim.holy](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tim.holy/32/52_2.png) [@tim.holy](https://discourse.julialang.org/u/tim.holy)\
**Post date:** [October 28, 2021, 9:55am UTC](https://discourse.julialang.org/t/indexing-an-arbitrary-array-over-its-last-dimension/70443/4 "2021-10-28T09:55:39Z")

</div>

> [@MattSainsbury-Dale](#):
>
> I was very surprised that splatting the colons from a `tuple` gives different results to splatting them from an `array` . I had thought that splatting involved throwing all elements of a container into the function, and so the container type doesn’t matter.

That is true, but there are nevertheless good reasons to use tuples whenever you can. Julia doesn’t know how many items a `Vector` contains, and so you’ll get inference failures if you index by splatting an array. In contrast, Julia does know how many items a tuple contains (if you’ve constructed it inferrably) and Life Will Be Good.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [October 28, 2021, 1:31pm UTC](https://discourse.julialang.org/t/indexing-an-arbitrary-array-over-its-last-dimension/70443/5 "2021-10-28T13:31:17Z")

</div>

Yes to tuples over arrays here. In addition, Zygote has many bugs, and you have to work around them to some degree… my guess is that this is related to [#599](https://github.com/FluxML/Zygote.jl/issues/599) — when something it splatted, it forgets that the thing wasn’t a tuple, and thus feeds a tuple to the gradient of `repeat`, which doesn’t like that. Probably.

> [@MattSainsbury-Dale](#):
>
> `selectdim()` is the neat function that I was looking for. However, it […] does some scalar indexing

That also sounds like a bug, for CUDA.jl.