Type inference with AbstractMatrix

Internals & Design
1

In the MixedModels package the bulk of the work fitting a model consists of updating a blocked Cholesky factor to new parameter values. The matrix being factored consists of a square set of blocks of possibly different types (e.g. SparseMatrixCSC, Diagonal, Matrix). All of the blocks have the same element type, usually Float64.

It appears that the way I am writing methods does not allow the compiler to infer the type of the return value from methods like

sqrtpwrss(m::LinearMixedModel{T}) where {T} = T(first(m.L.blocks[end, end]))

Applying @code_warntype returns

julia> @code_warntype MixedModels.sqrtpwrss(m1)
Variables
  #self#::Core.Compiler.Const(MixedModels.sqrtpwrss, false)
  m::LinearMixedModel{Float64}

Body::Any
1 ─ %1 = Base.getproperty(m, :L)::BlockArrays.BlockArray{Float64,2,R} where R<:AbstractArray{Float64,2}
│   %2 = Base.getproperty(%1, :blocks)::Array{R,2} where R<:AbstractArray{Float64,2}
│   %3 = Base.lastindex(%2, 1)::Int64
│   %4 = Base.lastindex(%2, 2)::Int64
│   %5 = Base.getindex(%2, %3, %4)::AbstractArray{Float64,2}
│   %6 = MixedModels.first(%5)::Any
│   %7 = ($(Expr(:static_parameter, 1)))(%6)::Any
└──      return %7

Is there a way I could write this to aid the type inference to produce a concrete type? Functions like this are called a lot.

2

Not really, but are you sure this is a performance problem? It feels like most of the computationally expensive functions should be behind function barriers and does be specialized on the type that gets passed in anyway.

3

Wouldn’t just a ::T be enough? Or do you want to avoid writing that manually?

1 Like