# Index into N-dimensional array with N indices

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?

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

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
``````

@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!

@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.)

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

1 Like