How to perform a reduction on wrappers of StaticArrays?

General Usage
,
1

The axes of a reduced array are computed using Base.reduced_indices. There seems to be a fundamental limitation here, that Base.reduced_indices expects the axis types to remain unchanged. However, given that the size of an axis is inbuilt into the axis type for a StaticArray, this is impossible. This leads to bugs such as:

julia> s = SArray{Tuple{2,2},Int,2,4}((1,2,3,4))
2×2 SMatrix{2, 2, Int64, 4} with indices SOneTo(2)×SOneTo(2):
 1  3
 2  4

julia> ss = StructArray{Complex{Int}}((s,s))
2×2 StructArray(::SMatrix{2, 2, Int64, 4}, ::SMatrix{2, 2, Int64, 4}) with eltype Complex{Int64} with indices SOneTo(2)×SOneTo(2):
 1+1im  3+3im
 2+2im  4+4im

julia> sum(ss, dims=1)
ERROR: ArgumentError: No method is implemented for reducing index range of type SOneTo{2}. Please implement
reduced_index for this index type or report this as an issue.

Stacktrace:
  [1] reduced_index(i::SOneTo{2})
    @ Base ./reducedim.jl:8
  [2] reduced_indices(inds::Tuple{SOneTo{2}, SOneTo{2}}, d::Int64)
    @ Base ./reducedim.jl:23
  [3] reduced_indices
    @ ./reducedim.jl:15 [inlined]
  [4] reducedim_initarray(A::StructArray{Complex{Int64}, 2, NamedTuple{(:re, :im), Tuple{SMatrix{2, 2, Int64, 4}, SMatrix{2, 2, Int64, 4}}}, Int64}, region::Int64, init::Complex{Int64}, #unused#::Type{Complex{Int64}})
    @ Base ./reducedim.jl:91
  [5] reducedim_initarray
    @ ./reducedim.jl:92 [inlined]
  [6] reducedim_init
    @ ./reducedim.jl:219 [inlined]
  [7] _mapreduce_dim
    @ ./reducedim.jl:371 [inlined]
  [8] #mapreduce#758
    @ ./reducedim.jl:357 [inlined]
  [9] #_sum#792
    @ ./reducedim.jl:1023 [inlined]
 [10] _sum
    @ ./reducedim.jl:1023 [inlined]
 [11] #_sum#791
    @ ./reducedim.jl:1022 [inlined]
 [12] _sum
    @ ./reducedim.jl:1022 [inlined]
 [13] #sum#765
    @ ./reducedim.jl:994 [inlined]
 [14] top-level scope
    @ REPL[7]:1

Note that StaticArrays don’t suffer from this limitation, as reduction methods are specialized for them. As a consequence, this works:

julia> sum(s, dims=1)
1×2 SMatrix{1, 2, Int64, 2} with indices SOneTo(1)×SOneTo(2):
 3  7

But the moment one deals with a wrapper of StaticArrays, this breaks down, as we hit generic Base fallbacks that rely on reduced_indices. What’s the best way to get this to work across the ecosystem, so that such a reduction operation works for any AbstractArray type that wraps a StaticArray?

One solution might be to define reduced_indices for Tuple{Vararg{SOneTo}}, but this can’t be done in a type-stable manner if we are to retain the size information. One way around that would be to just get rid of the sized axes, and convert them all to Base.OneTo. I wonder if there’s a better solution.

This still doesn’t solve the case of wrapper axes, eg. in OffsetArrays, as we’re dispatching on SOneTo. But at least it’s halfway there.