Base.stack of staticarrays creates non-static array

[wsmoses]

x/ref https://github.com/EnzymeAD/Enzyme.jl/issues/1714

[mcabbott]

Not so clear what the input is on that Enzyme issue, but the output is a Matrix.

But xref https://github.com/JuliaLang/julia/issues/52590, where the output is an MMatrix:

julia> stack([SA[1,2], SA[3,4], SA[5,6]])
2×3 Matrix{Int64}:
 1  3  5
 2  4  6

julia> stack(SA[SA[1,2], SA[3,4], SA[5,6]])
2×3 MMatrix{2, 3, Int64, 6} with indices SOneTo(2)×SOneTo(3):
 1  3  5
 2  4  6

[gdalle]

@mcabbott how would you go about handling the dims argument when creating SArray specializations for stack based on reduce(hcat/vcat)? When dims=1 or dims=2 I know what to do but other than that it seems less obvious.

[mcabbott]

I'm not so sure what you mean. Can you specify what the input is?

All legal values of dims do already work, e.g. with the input above:

julia> stack(SA[SA[1,2], SA[3,4], SA[5,6]]; dims=1)
3×2 MMatrix{3, 2, Int64, 6} with indices SOneTo(3)×SOneTo(2):
 1  2
 3  4
 5  6

julia> stack(SA[SA[1,2], SA[3,4], SA[5,6]]; dims=2)
2×3 MMatrix{2, 3, Int64, 6} with indices SOneTo(2)×SOneTo(3):
 1  3  5
 2  4  6

julia> stack(SA[SA[1,2], SA[3,4], SA[5,6]]; dims=3)
ERROR: ArgumentError: cannot stack slices ndims(x) = 1 along dims = 3

To make this return an SArray not MArray, it may make sense to just use Base's code path & convert the object before it's returned? But note that in any case it cannot be type-stable, since MMatrix{3, 2, ...} != MMatrix{2, 3, ...}.

Edit: Without dims it can be type-stable. Likely any implementation should overload internal functions, like _stack(dims::Colon, xs::Static...)

[gdalle]

To make this return an SArray not MArray

Yeah that would be the goal

But note that in any case it cannot be type-stable

I hadn't thought of that... @wsmoses do you think it is better to have stack be type-stable and return an Array, or type-unstable and return an SArray?

[mateuszbaran]

I hadn't thought of that... @wsmoses do you think it is better to have stack be type-stable and return an Array, or type-unstable and return an SArray?

The default choice in StaticArrays.jl is to return Array when an operation can't be type-stable with SArray. If someone wanted to suggest the other way, it'd have to be supported by some benchmarks.

[gdalle]

Yeah I hadn't realized that stack carries dimensionality information in the values of the kwarg dims, so it has to be type-unstable. I think defaulting on Array is a reasonable choice then.

[mcabbott]

Note that dims is optional, and cases without it can be type-stable (when the input has enough static axes). At present these give MMatrix, as that's what falls out of Base calling similar.

Note also that many examples above with dims are at present type-unstable because they make an MMatrix, not an Array.