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.