Type stability of vector of matrices

I’m puzzled by why the following function is type unstable:

function unstable(d, T)
    return [zeros(T,d,d) for _ = 1:d]
end

I got this by debugging a much more complicated function that had terrible performance; the reason turned out to be that the compiler couldn’t infer the type of this vector for some reason. I managed to create a type stable version

function stable(d, T)
    R = Vector{Matrix{T}}(undef, d)
    for i = 1:d
        R[i] = zeros(T,d,d)
    end
    return R
end

but that’s hideous, I’d rather avoid that if I can. Another possibility is

function also_stable(d, T)
    return zeros(T,d,d,d)
end

which is much nicer, but still, I’d rather work with a vector of matrices.

See Be aware of when Julia avoids specializing in the manual.

In particular, if you want the code to specialize on a type argument, you often need to be more explicit:

function stabilized(d, ::Type{T}) where T
    return [zeros(T,d,d) for _ = 1:d]
end

This code doesn’t run, because R = Vector{Matrix{T}} assigns R to the type itself, not an instance of the type. I guess you mean R = Vector{Matrix{T}}(undef, d)

(I’m not sure why this and the other example you showed are type stable … I wouldn’t count on this unless you use the ::Type{T} trick. Maybe it depends on whether the function gets inlined, which is decided heuristically by the compiler?)

5 Likes

Thanks a lot, that solves it. So it turns out that the behaviour I thought was anomalous is the standard, and the behaviour I thought was standard was anomalous. You’re right about the stable function, I edited it for clarity.

Your solution does generate another problem, though. In my code T was actually a keyword argument, with the default value Float64, and I don’t see how to parametrize that:

function unstable(d; T::Type=Float64)
    return [zeros(T,d,d) for _ = 1:d]
end

Any ideas?

Keyword arguments don’t participate in Julia’s dispatch rules, so the typical strategy is to pass off to an auxiliary positional function for dispatch, e.g. this is type stable:

@inline stable(d; T::Type=Float64) = _stable(d, T)
_stable(d, ::Type{T}) where {T} = [zeros(T,d,d) for _ = 1:d]

It is still unstable (infers an abstract type) if you use the T keyword in global scope, e.g. I get:

julia> @inferred stable(3; T=Int)
ERROR: return type Vector{Matrix{Int64}} does not match inferred return type Vector

but I put it into a function then the inlining works its magic:

julia> foo(d) = stable(d; T=Int)
foo (generic function with 1 method)

julia> @inferred foo(3)
3-element Vector{Matrix{Int64}}:
 [0 0 0; 0 0 0; 0 0 0]
 [0 0 0; 0 0 0; 0 0 0]
 [0 0 0; 0 0 0; 0 0 0]

Most commonly, you should avoid passing a type explicitly — if possible, you should compute the desired type from the other arguments.

If you must pass a type, it is a common convention to pass it as the first argument (as in zeros), e.g.:

foo(::Type{T}, x) where {T} = ....
foo(x) = foo(Float64, x) # default to Float64
3 Likes

I see, thanks. Indeed I always compute the type where I can, but in several functions I’m constructing a mathematical object (mutually unbiased bases in this particular case), and the desired type must be specified by the user.

I’m aware of the design of zeros and friends, and I’m not a fan. I wanted to do something more elegant. But the magic you found to make the keyword arguments work is even more clunky, so I’m afraid I’ll revert to respecting the convention.