# Selectdim behaviour for arrays of dimension greater than 2

**URL:** <https://discourse.julialang.org/t/selectdim-behaviour-for-arrays-of-dimension-greater-than-2/52211>\
**Category:** Performance\
**Created:** [December 22, 2020, 1:00am UTC](https://discourse.julialang.org/t/selectdim-behaviour-for-arrays-of-dimension-greater-than-2/52211 "2020-12-22T01:00:36Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![OlivierHnt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/olivierhnt/32/6227_2.png) [@OlivierHnt](https://discourse.julialang.org/u/OlivierHnt)\
**Post date:** [December 22, 2020, 1:00am UTC](https://discourse.julialang.org/t/selectdim-behaviour-for-arrays-of-dimension-greater-than-2/52211/1 "2020-12-22T01:00:36Z")

</div>

I just noticed a type-instability with `selectdim`. Consider the function

```nohighlight
f(A) = selectdim(A, 1, 1)

```

For a matrix, we have that the return type of `f` is a 1D `SubArray` (as expected):

```nohighlight
julia> A = rand(2,2);

julia> @code_warntype f(A)
Variables
  #self#::Core.Compiler.Const(f, false)
  A::Array{Float64,2}

Body::SubArray{Float64,1,Array{Float64,2},Tuple{Int64,Base.Slice{Base.OneTo{Int64}}},true}
1 ─ %1 = Main.selectdim(A, 1, 1)::SubArray{Float64,1,Array{Float64,2},Tuple{Int64,Base.Slice{Base.OneTo{Int64}}},true}
└── return %1

```

However, if we call `f` on an array of dimension `3`, the return type of `f` is a union of a 1D `SubArray` and a 2D `SubArray`:

```nohighlight
julia> A = rand(2,2,2);

julia> @code_warntype f(A)
Variables
  #self#::Core.Compiler.Const(f, false)
  A::Array{Float64,3}

Body::Union{SubArray{Float64,2,Array{Float64,3},Tuple{Int64,Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}}},true}, SubArray{Float64,1,Array{Float64,3},Tuple{Int64,Base.Slice{Base.OneTo{Int64}},Int64},true}}
1 ─ %1 = Main.selectdim(A, 1, 1)::Union{SubArray{Float64,2,Array{Float64,3},Tuple{Int64,Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}}},true}, SubArray{Float64,1,Array{Float64,3},Tuple{Int64,Base.Slice{Base.OneTo{Int64}},Int64},true}}
└── return %1

```

I expected the return type of `f` to b a 2D `SubArray` (hence be type stable) which would match the behaviour of `@views A[1,:,:]` no?  
Am I missing something here? when would `f` ever be a 1D `SubArray`?

---

<div class="post-metadata">

**Author:** ![mbauman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mbauman/32/31082_2.png) [@mbauman](https://discourse.julialang.org/u/mbauman)\
**Post date:** [December 22, 2020, 4:59am UTC](https://discourse.julialang.org/t/selectdim-behaviour-for-arrays-of-dimension-greater-than-2/52211/2 "2020-12-22T04:59:41Z")

</div>

> return type of `f` to b a 2D `SubArray` (hence be type stable)

`SubArray` also (necessarily) keeps track of the indices you used as part of the type:

```julia
julia> typeof(selectdim(A, 1, 1))
SubArray{Float64, 2, Array{Float64, 3}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}, true}

julia> typeof(selectdim(A, 2, 1))
SubArray{Float64, 2, Array{Float64, 3}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64, Base.Slice{Base.OneTo{Int64}}}, false}

julia> typeof(selectdim(A, 1, 1)) == typeof(selectdim(A, 2, 1))
false

```

Now, that said, it feels like we should be “saved” by constant propagation with the hard-coded dimension in `f(A) = selectdim(A, 1, 1)` — and indeed we _almost_ are! To get our list of indices, we recursively walk through each dimension, ask if it’s equal to the selected one, and conditionally use an int or a slice. The compiler happens to propagate constants through the first two recursions — two dimensions! — but then gives up. That is, it statically figured out that the indices in our SubArray are going to _definitely_ be a `1` and then a `:`… and then it gave up. That third dimension could either be a `1` or a `:`, depending upon the value of `d`. Yes, it feels funny, because the compiler knew the value of `d` in the first two dimensions, but now has forgotten.

Type stability isn’t about what it will actually return, it’s what it could _prove_ it _might_ return.

---

<div class="post-metadata">

**Author:** ![OlivierHnt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/olivierhnt/32/6227_2.png) [@OlivierHnt](https://discourse.julialang.org/u/OlivierHnt)\
**Post date:** [December 22, 2020, 6:43am UTC](https://discourse.julialang.org/t/selectdim-behaviour-for-arrays-of-dimension-greater-than-2/52211/3 "2020-12-22T06:43:56Z")

</div>

Absolutely 🙂 I thought constant propagation would yield type stability.

What feels funny to me is that in the example I gave, the compiler hesitates between a 1D `SubArray` (`SubArray{Float64,1,(...)}`) or a 2D `SubArray` (`SubArray{Float64,2,(...)}`) when it seems that it should only be concerned with subtypes of `SubArray{Float64,2,(...)}`.

Then, the follow-up question is why did the compiler give up after two dimensions?  
Is there a workaround?

---

<div class="post-metadata">

**Author:** ![mbauman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mbauman/32/31082_2.png) [@mbauman](https://discourse.julialang.org/u/mbauman)\
**Post date:** [December 22, 2020, 8:04am UTC](https://discourse.julialang.org/t/selectdim-behaviour-for-arrays-of-dimension-greater-than-2/52211/4 "2020-12-22T08:04:47Z")

</div>

The rubber hits the road here:

[https://github.com/JuliaLang/julia/blob/c708ca20cadc4c9e54edd643ae92f2fbf8bdc131/base/tuple.jl#L57](https://github.com/JuliaLang/julia/blob/c708ca20cadc4c9e54edd643ae92f2fbf8bdc131/base/tuple.jl#L57)

This is a recursive function that takes a value `v`, an index `i`, and a varargs list of arguments (the splatted axes) and will swap out the element at index `i` for value `v`. We start off with `i` being a literal constant — but as we progress, we do arithmetic on it… and at some point the compiler gives up (perhaps intentionally hitting a heuristic that’s trading effort — compile time — for return). Typically, the best tool for working with tuples is `ntuple` — that has the most smarts baked into it these days I think. We can try hacking in a replacement here:

```julia
julia> @inline function Base._setindex(v, i::Int, args...)
           ntuple(dim->ifelse(i==dim, v, args[dim]), length(args))
       end

julia> f(A)
2×2 view(::Array{Float64, 3}, 1, :, :) with eltype Float64:
 0.901717 0.427261
 0.130018 0.669248

julia> @code_warntype f(A)
Variables
  #self#::Core.Const(f)
  A::Array{Float64, 3}

Body::SubArray{Float64, 2, Array{Float64, 3}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}, true}

```

This is probably something we should propose doing…

---

<div class="post-metadata">

**Author:** ![OlivierHnt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/olivierhnt/32/6227_2.png) [@OlivierHnt](https://discourse.julialang.org/u/OlivierHnt)\
**Post date:** [December 22, 2020, 8:24am UTC](https://discourse.julialang.org/t/selectdim-behaviour-for-arrays-of-dimension-greater-than-2/52211/5 "2020-12-22T08:24:18Z")

</div>

Indeed, your fix looks very neat! 🙂 I’ll tag your reply as the solution for this post. Thanks for your help!

---

<div class="post-metadata">

**Author:** ![OlivierHnt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/olivierhnt/32/6227_2.png) [@OlivierHnt](https://discourse.julialang.org/u/OlivierHnt)\
**Post date:** [June 1, 2021, 2:27pm UTC](https://discourse.julialang.org/t/selectdim-behaviour-for-arrays-of-dimension-greater-than-2/52211/6 "2021-06-01T14:27:45Z")

</div>

Just as a follow-up: I opened [this PR](https://github.com/JuliaLang/julia/pull/40635) on GitHub. It is a faithful transcription of @mbauman’s post just above (cf. [the original issue](https://github.com/JuliaLang/julia/issues/39944) opened on GitHub and linked to the PR).

Additionally, the PR would really benefit from reviews 🙂.

EDIT: to add some motivation, here is a small showcase of the performance improvement:

```julia
# julia 1.6.1

using BenchmarkTools

f(A) = selectdim(A, 1, 1)

A = rand(10, 10, 10);

# @code_warntype f(A)
@btime f($A); # 148.503 ns (4 allocations: 128 bytes)

# implement the (one line) PR
Base._setindex(v, i::Int, args...) = ntuple(dim -> ifelse(i == dim, v, args[dim]), length(args))

# @code_warntype f(A)
@btime f($A); # 19.774 ns (0 allocations: 0 bytes)

```