Index into N-dimensional array with N indices

General Usage
1

Given an N-dimensional (custom) array A <: AbstractArray{T, N}, I want to index into this with N indices:

eg. if A <: AbstractArray{T, 2} I want A[:, :] or if A <: AbstractArray{T, 3} I want A[:, :, :]

This is a custom array that is strictly cartesian, so simply doing A[:] for linear indexing won’t work.

How do I programmatically do this? I’m assuming I need some metaprogramming here?

2

If all indexers are :, you’d just use copy(A) instead of A[:, :, ...].

However, as you asked this question, I’d suppose it was not an XY problem and what you expected could be the way to programmatically create indexing for variable dimensional arrays.

A zero-cost way is to use ntuple.

Given N is the type parameter comes from where clauses like,

function myfunc(A::TArray) where {N, T, TArray <: AbstractArray{T, N}}

We could use ntuple to generate the code we want:

# e.g., `ntuple(i -> :, 3)` creates a **constant** `(:, :, :)`

A[ntuple(dim  -> :, N)...]

Let us make this example more representative, for instance, index 1 for the first dimension, and : for the left:

A[ntuple(dim -> dim == 1 ? 1 : (:), N)...]

The generation is static, which can be seen from the generated code:

f(A::AbstractArray{T, N}) where {T, N} = A[ntuple(dim -> dim == 1 ? 1 : (:), N)...]
g(A) = A[1, :, :]

A = rand(3, 3, 3)

We can find the generated code is identity:

@code_typed g(A) 
julia> @code_typed g(A)
CodeInfo(
1 ─       Base.arraysize(A, 1)::Int64
│   %2  = Base.arraysize(A, 2)::Int64
│   %3  = Base.arraysize(A, 3)::Int64
│   %4  = Base.slt_int(%2, 0)::Bool
│   %5  = Base.ifelse(%4, 0, %2)::Int64
│   %6  = %new(Base.OneTo{Int64}, %5)::Base.OneTo{Int64}
│   %7  = Base.slt_int(%3, 0)::Bool
│   %8  = Base.ifelse(%7, 0, %3)::Int64
│   %9  = %new(Base.OneTo{Int64}, %8)::Base.OneTo{Int64}
│   %10 = %new(Base.Slice{Base.OneTo{Int64}}, %6)::Base.Slice{Base.OneTo{Int64}}
│   %11 = %new(Base.Slice{Base.OneTo{Int64}}, %9)::Base.Slice{Base.OneTo{Int64}}
└──       goto #6 if not true
2 ─ %13 = Core.tuple(1, %10, %11)::Tuple{Int64, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}
│   %14 = Base.arraysize(A, 1)::Int64
│         Base.arraysize(A, 2)::Int64
│         Base.arraysize(A, 3)::Int64
│   %17 = Base.slt_int(%14, 0)::Bool
│   %18 = Base.ifelse(%17, 0, %14)::Int64
│   %19 = Base.sle_int(1, 1)::Bool
│   %20 = Base.sle_int(1, %18)::Bool
│   %21 = Base.and_int(%19, %20)::Bool
│   %22 = Base.and_int(%21, true)::Bool
└──       goto #4 if not %22
3 ─       Base.nothing::Nothing
└──       goto #5
4 ─       invoke Base.throw_boundserror(A::Array{Float64, 3}, %13::Tuple{Int64, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}})::Union{}
└──       unreachable
5 ─       nothing::Nothing
6 ┄ %29 = invoke Base._unsafe_getindex($(QuoteNode(IndexLinear()))::IndexLinear, A::Array{Float64, 3}, 1::Int64, %10::Base.Slice{Base.OneTo{Int64}}, %11::Base.Slice{Base.OneTo{Int64}})::Matrix{Float64}
└──       goto #7
7 ─       goto #8
8 ─       return %29
) => Matrix{Float64}
@code_typed f(A) 
julia> @code_typed f(A)
CodeInfo(
1 ─       Base.arraysize(A, 1)::Int64
│   %2  = Base.arraysize(A, 2)::Int64
│   %3  = Base.arraysize(A, 3)::Int64
│   %4  = Base.slt_int(%2, 0)::Bool
│   %5  = Base.ifelse(%4, 0, %2)::Int64
│   %6  = %new(Base.OneTo{Int64}, %5)::Base.OneTo{Int64}
│   %7  = Base.slt_int(%3, 0)::Bool
│   %8  = Base.ifelse(%7, 0, %3)::Int64
│   %9  = %new(Base.OneTo{Int64}, %8)::Base.OneTo{Int64}
│   %10 = %new(Base.Slice{Base.OneTo{Int64}}, %6)::Base.Slice{Base.OneTo{Int64}}
│   %11 = %new(Base.Slice{Base.OneTo{Int64}}, %9)::Base.Slice{Base.OneTo{Int64}}
└──       goto #6 if not true
2 ─ %13 = Core.tuple(1, %10, %11)::Tuple{Int64, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}
│   %14 = Base.arraysize(A, 1)::Int64
│         Base.arraysize(A, 2)::Int64
│         Base.arraysize(A, 3)::Int64
│   %17 = Base.slt_int(%14, 0)::Bool
│   %18 = Base.ifelse(%17, 0, %14)::Int64
│   %19 = Base.sle_int(1, 1)::Bool
│   %20 = Base.sle_int(1, %18)::Bool
│   %21 = Base.and_int(%19, %20)::Bool
│   %22 = Base.and_int(%21, true)::Bool
└──       goto #4 if not %22
3 ─       Base.nothing::Nothing
└──       goto #5
4 ─       invoke Base.throw_boundserror(A::Array{Float64, 3}, %13::Tuple{Int64, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}})::Union{}
└──       unreachable
5 ─       nothing::Nothing
6 ┄ %29 = invoke Base._unsafe_getindex($(QuoteNode(IndexLinear()))::IndexLinear, A::Array{Float64, 3}, 1::Int64, %10::Base.Slice{Base.OneTo{Int64}}, %11::Base.Slice{Base.OneTo{Int64}})::Matrix{Float64}
└──       goto #7
7 ─       goto #8
8 ─       return %29
) => Matrix{Float64}
2 Likes
3

To answer my own question, use Base.Cartesian · The Julia Language

e.g.:

@generated function test(A::Array{T, N}) where {T, N}
   quote
       Base.Cartesian.@nref $N A j->(i_j = Colon())
   end
end
4

@thautwarm Appreciate your answer here - I did not know about the ntuple(f, N) function - that’s a nice tidy way to do this too!

5

@thautwarm In fact, I’m marking yours as the solution because it’s much simpler and elegant than some Base.Cartesian.

(As an aside, yes, copy() would normally be the right way to do this, but in actual fact this is to implement copy() for a custom array.)

6

You may look at GitHub - SciML/EllipsisNotation.jl: Implements the notation `..` for indexing arrays. as well

1 Like