Index into N-dimensional array with N indices

General Usage
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