Indexing an arbitrary array over its last dimension

General Usage
,
1

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:

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:

# 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:

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?

2 Likes
2

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

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.

3 Likes
3

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.

4

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.

5

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 — 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.

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