Feature request: concatenation may preserve static sizes

[jishnub]

Currently, concatenation passes the sizes of the components to determine the type of the container. It might be possible to use the axes instead, thereby preserving static sizes. As an example:

julia> using StaticArrays, FillArrays

julia> a = SMatrix{2,2}(1:4)
2×2 SMatrix{2, 2, Int64, 4} with indices SOneTo(2)×SOneTo(2):
 1  3
 2  4

julia> row = Ones(SOneTo(1), axes(a,2))
1×2 Ones{Float64, 2, Tuple{SOneTo{1}, SOneTo{2}}} with indices SOneTo(1)×SOneTo(2)

julia> [a; row]
3×2 Matrix{Float64}:
 1.0  3.0
 2.0  4.0
 1.0  1.0

Since all components have their sizes statically known, it should be possible to preserve that information in the container. This will perhaps require some function that accepts SOneTo(2) and SOneTo(1) as arguments and returns SOneTo(3). I'm uncertain if there is an existing function that does this. However, once this is available, one may pass the result to similar to generate the container.

[Moelf]

what has this to do with Julia base? it should be an issue opened in StaticArrays?

julia> Meta.@lower [a; row]
:($(Expr(:thunk, CodeInfo(
    @ none within `top-level scope`
1 ─ %1 = Base.vcat(a, row)
└──      return %1
))))

they need to specialize vcat on their types

[jishnub]

This isn't specializing on array types, but on axes. The arrays may come from different packages, but they share a common feature: their axes are statically sized (originating from StaticArrays). Currently, vcat in Base immediately converts axes to sizes, and uses sizes subsequently in the concatenation. The proposal here is to split the allocation of the output container into a separate function, one that uses the axes. The fallback definition in Base would proceed as it does now, by converting the axes to sizes and allocating the container. This function may then be specialized by StaticArrays for their axis type.

In short, the proposal here is to add a function that may then be specialized for axes, not arrays. This function doesn't exist currently, IIUC. vcat is already specialized for static arrays, but there is no notion of specializing vcat for statically sized axes.

The current workaround is to wrap row in a SizedArray so that it becomes a static array, and can use the specialized vcat.

julia> [a; row]
3×2 Matrix{Float64}:
 1.0  3.0
 2.0  4.0
 1.0  1.0

julia> [a; SizedArray{Tuple{size(row)...}}(row)]
3×2 SMatrix{3, 2, Float64, 6} with indices SOneTo(3)×SOneTo(2):
 1.0  3.0
 2.0  4.0
 1.0  1.0

With this proposal, [a; row] would also be a static array.

[Moelf]

With this proposal, [a; row] would also be a static array.

what this goes to ultimately depends on if Base.vcat , how would Julia base know what is StaticArray which comes from a package and how Julia base know what "promotion rule" should it adopt?

[jishnub]

Base doesn't need to know about StaticArrays. It only needs to define a function like this:

julia> vcat_container(a, ax1::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}}, ax2::AbstractUnitRange) = similar(a, sum(length, ax1), length(ax2))
vcat_container (generic function with 1 method)

julia> vcat_container(As::Vararg{AbstractArray}) = vcat_container(As[1], map(Base.axes1, As), axes(As[1],2))
vcat_container (generic function with 2 methods)

StaticArrays may specialize this as

julia> vcat_container(a, ax1::Tuple{SOneTo,Vararg{SOneTo}}, ax2::SOneTo) = similar(a, SOneTo(sum(length, ax1)), ax2)
vcat_container (generic function with 3 methods)

With this defined, we obtain

julia> vcat_container(a, row)
3×2 MMatrix{3, 2, Int64, 6} with indices SOneTo(3)×SOneTo(2):
 140609240823168   1
               1   8
          559109  -1

which is statically sized. I haven't been careful with the eltype here, but that shouldn't be hard to incorporate.

[mcabbott]

This is related to handling offset arrays, #37629

IMO the cat functions should all work through axes not size, and either

• use the existing axis when perpendicular to concatenation, e.g. axes(hcat(v, v), 1) == axes(v, 1). This is what stack almost always does.
• use some new axis_cat function when the axes are concatenated, e.g. axes(vcat(v, v), 1) == axis_cat(axes(v,1), axes(v,1)).

This axis_cat would default to Base.OneTo(length + length). It could be overloaded by StaticArrays. Maybe some thought should be given to allowing it to be extended by several packages while avoiding ambiguities.

I think your vcat_container(a, ax1, ax2) is just similar(a, axis_cat(ax1, ax2)).

(The stack PR also added a Base._vec_axis, which returns OneTo of the product of lengths. This only matters for stack(multi_dim_container; dims=1) == stack(vec(multi_dim_container); dims=1), which can't have all the original axes.)