`stack` errors when inner or outer array uses custom indices

[sethaxen]

When arrays that use custom index types as suggested in the docs are passed to stack, everything works if the wrapped array uses the same custom index type. However, if the inner and outer arrays have different index types, an error is raised on this line:

https://github.com/JuliaLang/julia/blob/ce292c1ba4f5ec771479228169383f6378d35d45/base/abstractarray.jl#L2792

This calls the following similar fallback, which replaces Base.OneTo axes with Ints, creating a mixture of Ints and custom indices for which no similar method is defined:

https://github.com/JuliaLang/julia/blob/ce292c1ba4f5ec771479228169383f6378d35d45/base/abstractarray.jl#L839-L844

(it seems that following the instructions in the comment would both be type-piracy and create ambiguities)

Here's an minimal example where we define a custom array type that uses custom indices as defined in CustomUnitRanges.jl and then call stack on these arrays:

module MyArrayType

# use URange defined in CustomUnitRanges, recommended in Julia docs
# https://github.com/JuliaArrays/CustomUnitRanges.jl#usage
using CustomUnitRanges: filename_for_urange
include(filename_for_urange)

# array type that will have custom axes
struct MyArray{T,N} <: AbstractArray{T,N}
    data::Array{T,N}
end

Base.parent(A::MyArray) = A.data

# iterator interface
# https://docs.julialang.org/en/v1/manual/interfaces/#man-interface-iteration
Base.iterate(A::MyArray) = iterate(parent(A))
Base.iterate(A::MyArray, state) = iterate(parent(A), state)

# array interface
# https://docs.julialang.org/en/v1/manual/interfaces/#man-interface-array
Base.size(A::MyArray) = size(parent(A))
Base.getindex(A::MyArray, I::Int...) = getindex(parent(A), I...)
Base.setindex!(A::MyArray, v, I::Int...) = setindex!(parent(A), v, I...)

# custom indexing interface
# https://docs.julialang.org/en/v1/devdocs/offset-arrays/#Writing-custom-array-types-with-non-1-indexing
Base.axes(A::MyArray) = map(Base.Slice ∘ Base.Fix1(URange, 1), size(parent(A)))
function Base.similar(
    A::AbstractArray, ::Type{T},
    dims::Tuple{Base.Slice{<:URange},Vararg{Base.Slice{<:URange}}},
) where {T}
    return MyArray(similar(A, T, map(length, dims)))
end

end

julia> A1 = MyArrayType.MyArray([MyArrayType.MyArray(randn(4)) for _ in 1:3, _ in 1:2]);

julia> stack(A1)
4×3×2 Main.MyArrayType.MyArray{Float64, 3} with indices URange(1,4)×URange(1,3)×URange(1,2):
[:, :, 1] =
  0.39151   -1.02691   -0.399949
 -1.81154    0.832518  -1.46895
  0.116189   0.108829  -0.0630371
  1.05827   -2.42168    0.143143

[:, :, 2] =
 1.21023     0.442859  -0.755589
 0.0321352  -0.339096   1.92049
 1.03333     0.667389  -1.94571
 0.829069    0.999837  -0.728966

julia> A2 = MyArrayType.MyArray([randn(4) for _ in 1:3, _ in 1:2]);

julia> stack(A2)
ERROR: MethodError: no method matching similar(::Vector{Float64}, ::Type{Float64}, ::Tuple{Int64, Base.Slice{Main.MyArrayType.URange{Int64}}, Base.Slice{Main.MyArrayType.URange{Int64}}})

Closest candidates are:
  similar(::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{N}}, ::Type{ElType}, ::Any) where {N, ElType}
   @ Base broadcast.jl:212
  similar(::Base.Broadcast.Broadcasted{Base.Broadcast.ArrayConflict}, ::Type{ElType}, ::Any) where ElType
   @ Base broadcast.jl:217
  similar(::Vector{T}, ::Type) where T
   @ Base array.jl:371
  ...

Stacktrace:
 [1] similar(::Vector{Float64}, ::Type{Float64}, ::Base.OneTo{Int64}, ::Base.Slice{Main.MyArrayType.URange{Int64}}, ::Base.Slice{Main.MyArrayType.URange{Int64}})
   @ Base ./abstractarray.jl:838
 [2] _typed_stack(::Colon, ::Type{Float64}, ::Type{Vector{Float64}}, A::Main.MyArrayType.MyArray{Vector{Float64}, 2}, Aax::Tuple{Base.Slice{Main.MyArrayType.URange{Int64}}, Base.Slice{Main.MyArrayType.URange{Int64}}})
   @ Base ./abstractarray.jl:2791
 [3] _typed_stack(::Colon, ::Type{Float64}, ::Type{Vector{Float64}}, A::Main.MyArrayType.MyArray{Vector{Float64}, 2})
   @ Base ./abstractarray.jl:2787
 [4] _stack(dims::Function, #unused#::Base.HasShape{2}, iter::Main.MyArrayType.MyArray{Vector{Float64}, 2})
   @ Base ./abstractarray.jl:2777
 [5] _stack(dims::Colon, iter::Main.MyArrayType.MyArray{Vector{Float64}, 2})
   @ Base ./abstractarray.jl:2769
 [6] stack(iter::Main.MyArrayType.MyArray{Vector{Float64}, 2}; dims::Function)
   @ Base ./abstractarray.jl:2737
 [7] stack(iter::Main.MyArrayType.MyArray{Vector{Float64}, 2})
   @ Base ./abstractarray.jl:2737
 [8] top-level scope
   @ REPL[47]:1

julia> A3 = [MyArrayType.MyArray(randn(4)) for _ in 1:3, _ in 1:2];

julia> stack(A3)
ERROR: MethodError: no method matching similar(::Main.MyArrayType.MyArray{Float64, 1}, ::Type{Float64}, ::Tuple{Base.Slice{Main.MyArrayType.URange{Int64}}, Int64, Int64})

Closest candidates are:
  similar(::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{N}}, ::Type{ElType}, ::Any) where {N, ElType}
   @ Base broadcast.jl:212
  similar(::Base.Broadcast.Broadcasted{Base.Broadcast.ArrayConflict}, ::Type{ElType}, ::Any) where ElType
   @ Base broadcast.jl:217
  similar(::AbstractArray, ::Type{T}, ::Tuple{Base.Slice{<:Main.MyArrayType.URange}, Vararg{Base.Slice{<:Main.MyArrayType.URange}}}) where T
   @ Main.MyArrayType REPL[30]:21
  ...

Stacktrace:
 [1] similar(::Main.MyArrayType.MyArray{Float64, 1}, ::Type{Float64}, ::Base.Slice{Main.MyArrayType.URange{Int64}}, ::Base.OneTo{Int64}, ::Base.OneTo{Int64})
   @ Base ./abstractarray.jl:838
 [2] _typed_stack(::Colon, ::Type{Float64}, ::Type{Main.MyArrayType.MyArray{Float64, 1}}, A::Matrix{Main.MyArrayType.MyArray{Float64, 1}}, Aax::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}})
   @ Base ./abstractarray.jl:2791
 [3] _typed_stack(::Colon, ::Type{Float64}, ::Type{Main.MyArrayType.MyArray{Float64, 1}}, A::Matrix{Main.MyArrayType.MyArray{Float64, 1}})
   @ Base ./abstractarray.jl:2787
 [4] _stack(dims::Function, #unused#::Base.HasShape{2}, iter::Matrix{Main.MyArrayType.MyArray{Float64, 1}})
   @ Base ./abstractarray.jl:2777
 [5] _stack(dims::Colon, iter::Matrix{Main.MyArrayType.MyArray{Float64, 1}})
   @ Base ./abstractarray.jl:2769
 [6] stack(iter::Matrix{Main.MyArrayType.MyArray{Float64, 1}}; dims::Function)
   @ Base ./abstractarray.jl:2737
 [7] stack(iter::Matrix{Main.MyArrayType.MyArray{Float64, 1}})
   @ Base ./abstractarray.jl:2737
 [8] top-level scope
   @ REPL[49]:1

julia> versioninfo()
Julia Version 1.9.0-beta4
Commit b75ddb787ff (2023-02-07 21:53 UTC)
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 8 × 11th Gen Intel(R) Core(TM) i5-1135G7 @ 2.40GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, tigerlake)
  Threads: 8 on 8 virtual cores

(jl_0uKWW2) pkg> st
Status `/tmp/jl_0uKWW2/Project.toml`
  [dc8bdbbb] CustomUnitRanges v1.0.2

[rafaqz]

As this functionality was seemingly designed mostly for OffsetArrays.jl: https://github.com/JuliaLang/julia/blob/ce292c1ba4f5ec771479228169383f6378d35d45/base/abstractarray.jl#L839-L844

it should be noted that this does not produce an error when OffsetArrays.jl is loaded and it's type pirated methods are available.

julia> using OffsetArrays

julia> stack(A2)
4×3×2 OffsetArray(::Array{Float64, 3}, 1:4, 1:3, 1:2) with eltype Float64 with indices 1:4×1:3×1:2:
[:, :, 1] =
  0.488586   -1.93562   -0.225797
 -0.307972   -1.7415     0.316024
  0.0702133   1.00626    1.00785
 -0.72046    -0.679983   0.00612942

[:, :, 2] =
 -1.59223    0.291592   1.48481
 -1.87154   -0.727295   0.378543
 -1.08153    1.35971    1.63377
  0.819348  -0.132509  -1.36885

julia> stack(A3)
4×3×2 OffsetArray(::Array{Float64, 3}, 1:4, 1:3, 1:2) with eltype Float64 with indices 1:4×1:3×1:2:
[:, :, 1] =
 -1.13547    0.695512  -0.0426222
  0.296079   1.85618   -0.247409
 -0.879178  -0.417559  -0.447225
 -0.791633   0.283552  -0.221949

[:, :, 2] =
 -0.76631    0.116022   1.31183
 -1.00218    2.7132     0.148179
  0.0976066  0.551973  -1.16329
 -0.966548   0.382295   0.341279

Meaning e.g. DimensionalData.jl tests on stack will pass because it tests against OffsetArrays.jl. But stack will not work in a fresh session without OffsetArrays.jl loaded.

[mcabbott]

stack seems like a distraction here.

With the above MyArray type, similar with a mix of its special axes and ordinary Base.OneTo fails:

julia> let x = MyArrayType.MyArray([1,2,3])
         similar(x, Int, axes(x)..., axes([4,5])...)  # dispatches to Base's mix-of-ranges-and-Int method
       end
ERROR: MethodError: no method matching similar(::Main.MyArrayType.MyArray{Int64, 1}, ::Type{Int64}, ::Tuple{Base.Slice{Main.MyArrayType.URange{Int64}}, Int64})
Stacktrace:
 [1] similar(::Main.MyArrayType.MyArray{Int64, 1}, ::Type{Int64}, ::Base.Slice{Main.MyArrayType.URange{Int64}}, ::Base.OneTo{Int64})
   @ Base ./abstractarray.jl:838
 [2] top-level scope
   @ REPL[4]:2

julia> using OffsetArrays  # this supplies the missing method, so exactly the same code now runs:

julia> let x = MyArrayType.MyArray([1,2,3])
         similar(x, Int, axes(x)..., axes([4,5])...)
       end
3×2 OffsetArray(::Matrix{Int64}, 1:3, 1:2) with eltype Int64 with indices 1:3×1:2:
 5648757792           0
          0  5648757760
 4780822496           0

If you want similar to notice your special axis type, I would have thought you ought to define methods accepting at least one of this special type, mixed with Base types. Rather the only covering the case when all axes are your special type. (Not sure why it's Slice{URange} here but that doesn't matter.) I see now that the docs recommend only the same-type similar.

Adding ever more methods to similar isn't a very extensible mechanism. It can't work if two packages both try it. Perhaps it could be replaced by something more like to_indices which peels them off one by one. Or perhaps it could use promotion, with a method to catch all mixed cases which calls similar(x,T,promote(axes...)...).