ArrayInterface.size vs StaticArrays.Size

[cscherrer]

Hi,

I'd like to get a better understanding of the intending idioms for working with arrays, especially the case where we want different methods for the static vs dynamic case.

StaticArrays.Size solves the problem like this:

julia> Size(zeros(SMatrix{3, 4}))
Size(3, 4)

julia> Size(zeros(3, 4))
Size(StaticArrays.Dynamic(), StaticArrays.Dynamic())

Though we might have some missing results, it's nice that the size computation itself is always static. Working with types is also convenient:

julia> Size(typeof(zeros(SMatrix{3, 4})))
Size(3, 4)

julia> Size(typeof(zeros(3, 4)))
Size(StaticArrays.Dynamic(), StaticArrays.Dynamic())

By contrast, here's ArrayInterface.size:

julia> ArrayInterface.size(zeros(SMatrix{3, 4}))
(static(3), static(4))

julia> ArrayInterface.size(zeros(3, 4))
(3, 4)

That last result might seem convenient, but computing this requires dynamic dispatch, which seems to defeat the purpose. And it fails on types:

julia> ArrayInterface.size(typeof(zeros(3, 4)))
ERROR: MethodError: no method matching size(::Type{Matrix{Float64}})
Closest candidates are:
  size(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:524
  size(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}, ::Integer) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:523
  size(::Union{LinearAlgebra.Cholesky, LinearAlgebra.CholeskyPivoted}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/cholesky.jl:442
  ...
Stacktrace:
 [1] axes
   @ ./abstractarray.jl:89 [inlined]
 [2] axes
   @ ~/.julia/packages/ArrayInterface/CmPZc/src/axes.jl:166 [inlined]
 [3] size(a::Type{Matrix{Float64}})
   @ ArrayInterface ~/.julia/packages/ArrayInterface/CmPZc/src/size.jl:18
 [4] top-level scope
   @ REPL[95]:1

julia> ArrayInterface.size(typeof(zeros(3, 4)))
ERROR: MethodError: no method matching size(::Type{Matrix{Float64}})
Closest candidates are:
  size(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:524
  size(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}, ::Integer) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/qr.jl:523
  size(::Union{LinearAlgebra.Cholesky, LinearAlgebra.CholeskyPivoted}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/cholesky.jl:442
  ...
Stacktrace:
 [1] axes
   @ ./abstractarray.jl:89 [inlined]
 [2] axes
   @ ~/.julia/packages/ArrayInterface/CmPZc/src/axes.jl:166 [inlined]
 [3] size(a::Type{Matrix{Float64}})
   @ ArrayInterface ~/.julia/packages/ArrayInterface/CmPZc/src/size.jl:18
 [4] top-level scope
   @ REPL[95]:1

Is ArrayInterface.size intended to be used for building generated functions or for something else? In what cases does it have benefits over StaticArrays.Size?

[Tokazama]

How does this require dynamic dispatch? All the statically known info is computed in the type domain and inlined so that it is known at compile time. If you want to operate on just the type you can use ArrayInterface.known_size.

For reference, here are some related methods:

• known_length(x)
• known_first(x)
• known_last(x)
• known_step(x)

[cscherrer]

Say I'm building a mul(A::AbstractMatrix, x::AbstractVector), and I dispatch using ArrayInterface.size. It seems to me this would only be entirely static when all sizes are known statically.

known_size seems to partially replicate Size, but the result is not type-level, so it still can't be used directly for dispatch.

I think I'm just missing some context on intended usage. Do you know of examples where these methods make it easier to build fast code than just using StaticArrays.Size?

[cscherrer]

Oh, maybe something like

struct Sized{S,T}
    x::T
end

Sized(x::T) where {T} = Sized{ArrayInterface.known_size(T), T}(x)

?

Then dispatch is easier:

julia> f(::Sized{(3,)}) = "three"
f (generic function with 1 method)

julia> f(x) = "not three"
f (generic function with 2 methods)

julia> x = Sized(SizedVector{3}([1,2,3]))
Sized{(3,), SizedVector{3, Int64, Vector{Int64}}}([1, 2, 3])

julia> f(x)
"three"

[Tokazama]

If you want a single type to dispatch on you could do VecSize{S} = Tuple{StaticInt{S}} and MatSize{S1,S2} = Tuple{StaticInt{S1},StaticInt{S2}}.

If you don't know how many dimensions it is and all you know is that you have a type T = Tuple{Vararg{StaticInt}}, then you can get the regular tuple with Static.known(T).

[cscherrer]

Ok, how would you decide whether to take this approach or to use StaticArrays.Size?

[Tokazama]

I wouldn't ever use StaticArrays.Size. I think some maintainers for StaticArrays.jl agree. I honestly don't see any benefit to Size.

[cscherrer]

Ah, there's the missing context. That's really helpful, thanks!

[Tokazama]

Of course! Happy to help

[chriselrod]

and I dispatch using ArrayInterface.size. It seems to me this would only be entirely static when all sizes are known statically.

Not sure what you meant here, but ArrayInterface.size shouldn't be resulting in dynamic dispatches.

julia> using StrideArrays, Static, ArrayInterface

julia> A = StrideArray{Float64}(undef, (static(2), 3, static(4), 5));

julia> ArrayInterface.size(A)
(static(2), 3, static(4), 5)

julia> typeof(ans)
Tuple{StaticInt{2}, Int64, StaticInt{4}, Int64}

julia> Static.known(typeof(ArrayInterface.size(A)))
(2, nothing, 4, nothing)

You can dispatch on a Tuple{StaticInt{2}, Int, StaticInt{4}, Int} just fine, and the type is known at compile time.
You may have to be careful indexing into it, i.e. the indexes should be known at compile time.

I'll echo Tokazama's points.

Also, map tends to be type stable when working with heterogenous tuples, while broadcast isn't.
My code working with tuples thus tends to be a mix of

1. map
2. first and Base.tail in recursive functions, where each call pops off the first remaining element, processes it, and then calls itself using Base.tail as the arg.
3. @generated

[cscherrer]

Not sure what you meant here, but ArrayInterface.size shouldn't be resulting in dynamic dispatches.

I think it depends how you use it - I'm just trying to figure out the right idioms.

Say you're building a mul(m::AbstractMatrix, y::Abstractvector). You might try something like

function mul(m::AbstractMatrix, v::Abstractvector)
    m_size = ArrayInterface.size(m)
    v_size = ArrayInterface.size(v)
    _mul(m, v, m_size, v_size)

where _mul has some generated methods for static sizes. It seemed to me if m or v have any dynamic sizes this would require dynamic dispatch. Maybe the compiler is smart enough to work around this?

Then, if you're building a generated function, ArrayInterface.size doesn't help, because it fails on types. It's helpful to know about known_size.

Basically what I'm looking for is a design pattern for writing functions that split into generic methods vs static methods with generated functions.

In the multivariate normal example we're looking at, I'm currently doing something like

logdensity(d::MvNormal{(:ω,), <:Tuple{<:Cholesky}}, y) = logdensity_mvnormal_ω(getfield(d.ω, :factors), y)

logdensity_mvnormal_ω(UL::Union{UpperTriangular, LowerTriangular}, y::AbstractVector) = logdensity_mvnormal_ω(KnownSize(UL), y)

where KnownSize is

struct KnownSize{S, T}
    value::T
end

KnownSize(x::T) where {T} = KnownSize{Tuple{ArrayInterface.known_size(T)...}, T}(x)

This gives a size wrapper around the value. Dispatch is easy, you can keep the same argument slots, and it's easy to unwrap. But maybe there's a cleaner way to do this?

[Tokazama]

It seemed to me if m or v have any dynamic sizes this would require dynamic dispatch. Maybe the compiler is smart enough to work around this?

It still shouldn't be dynamic dispatch (although this is helping me realize how obtuse some of the code currently is without more detailed documentation).

If you look at related tests, you can see that it's all fully inferred.

I would just call logdensity_mvnormal_ω(ArrayInterface.size(UL), y).
Then if you want a unique function for each combination of "staticness" define:

logdensity_mvnormal_ω(s::Tuple{Int,Int}, y) = ...
logdensity_mvnormal_ω(s::Tuple{StaticInt{S1},Int}, y) where {S1} = ...
logdensity_mvnormal_ω(s::Tuple{Int,StaticInt{S2}}, y) where {S2} = ...
logdensity_mvnormal_ω(s::Tuple{StaticInt{S1},StaticInt{S2}}, y) where {S1,S2} = ...

Depending on what you need to do you might want to make any of the last three @generated

You could also do something like

@generated function logdensity_mvnormal_ω(s::Tuple{S1,S2}, y) where {S1,S2}
    s1 = known(S1) === nothing ? :(getfield(s, 1)) : known(S1)
    s2 = known(S2) === nothing ? :(getfield(s, 2)) : known(S2)

    ...do stuff
end

[cscherrer]

Ok, this is making some more sense. Thank you for your help!

[chriselrod]

Then, if you're building a generated function, ArrayInterface.size doesn't help, because it fails on types. It's helpful to know about known_size.

Looking at your earlier example, something like this should also work:

using Static, ArrayInterface
function mul(m::AbstractMatrix, v::Abstractvector)
    m_size = ArrayInterface.size(m)
    v_size = ArrayInterface.size(v)
    _mul(m, v, m_size, v_size)
end
@generated function _mul(m, v, m_size::MS, v_size::VS) where {MS,VS}
    # `ms` and `vs` are tuples, with values for known information and `nothing` when the sizes were dynamic
    ms = known(MS)
    vs = known(VS)
    m_sizes = Any[]
    q = quote end
    # build the expression. If `ms[i] === nothing`, you'd use `:(m_size[$i])` for the size,
    # otherwise you can use `ms[i]` or manually unroll the expression by `ms[i]`.
end

I agree that this could use more documentation.

Ideas on how to make this more convenient are welcome.

And to repeat, we do want to keep all the ArrayInterface.size calls outside of generated functions to avoid possible world age issues.

[cscherrer]

I've thought about this some more while working on https://github.com/cscherrer/MultivariateMeasures.jl. There are two main cases that come up. First, we may have full control over the generic implementation. That's like this:

Looking at your earlier example, something like this should also work:

using Static, ArrayInterface
function mul(m::AbstractMatrix, v::Abstractvector)
    m_size = ArrayInterface.size(m)
    v_size = ArrayInterface.size(v)
    _mul(m, v, m_size, v_size)
end
@generated function _mul(m, v, m_size::MS, v_size::VS) where {MS,VS}
    # `ms` and `vs` are tuples, with values for known information and `nothing` when the sizes were dynamic
    ms = known(MS)
    vs = known(VS)
    m_sizes = Any[]
    q = quote end
    # build the expression. If `ms[i] === nothing`, you'd use `:(m_size[$i])` for the size,
    # otherwise you can use `ms[i]` or manually unroll the expression by `ms[i]`.
end

Let's say m is i by j, and v has length j. Some or all of these might be known statically. So maybe we have a strange case where

ArrayInterface.size(m) == (StaticInt(5), 4)
ArrayInterface.size(v) == (StaticInt(4),)

Making this fast should be no problem, but my earlier concern is that even asking for ArrayInterface.size(m) would (I thought) trigger a dynamic dispatch to get its column count.

The second situation is where a generic mul(m::AbstractMatrix, v::Abstractvector) already exists, and we need to add optimized methods for static cases. Is there a nice way to do this?