Closures over types are type unstable

The closure returned by the following function

function f(x::AbstractArray)
           T = eltype(x)
           ax = axes(x)
           return () -> zeros(T, 1000)
end

is type unstable:

@code_warntype f(rand(3))()
MethodInstance for (::var"#41#42"{DataType})()
  from (::var"#41#42")() in Main at REPL[37]:4
Arguments
  #self#::var"#41#42"{DataType}
Body::Vector
1 ─ %1 = Core.getfield(#self#, :T)::DataType
│   %2 = Main.zeros(%1, 1000)::Vector
└──      return %2

How can I modify it so that the function f returns a closure which is type stable?

This is an instance of the infamous closure bug

If I understand correctly, the issue is that your closure views T as something “global” that could possibly change after its definition.
Here is one way to solve it:

julia> function f(x::AbstractArray)
           T = eltype(x)
           ax = axes(x)
           g() = let T
               zeros(T, 1000)
           end
           return g
       end
f (generic function with 1 method)

julia> @code_warntype f(rand(3))()
MethodInstance for (::var"#g#11")()
  from (::var"#g#11")() in Main at /home/guillaume/Downloads/scratchpad.jl:4
Arguments
  #self#::Core.Const(var"#g#11"())
Locals
  T::Union{}
Body::Union{}
1 ─     Core.NewvarNode(:(T))
│       Main.zeros(T, 1000)
└──     Core.Const(:(return %2))

PR: Fix type instability of closures capturing types (2) by simeonschaub · Pull Request #40985 · JuliaLang/julia · GitHub

You could also do

julia> function f(x::AbstractArray{T}) where T
       ax = axes(x)
       return () -> zeros(T, 1000)
       end
f (generic function with 1 method)

julia> @code_warntype f(rand(3))()
MethodInstance for (::var"#1#2"{Float64})()
  from (::var"#1#2"{T})() where T in Main at REPL[1]:3
Static Parameters
  T = Float64
Arguments
  #self#::Core.Const(var"#1#2"{Float64}())
Body::Vector{Float64}
1 ─ %1 = Main.zeros($(Expr(:static_parameter, 1)), 1000)::Vector{Float64}
└──      return %1
1 Like

Do you happen to know about the timeline for merging that PR?

That infers as Union{}, because it’ll throw rather than return. But replacing zeros(T, 1000) with zeros(eltype(T), 1000) does infer on at least Julia 1.7, 1.8 and master.

I don’t know a timeline for merging the PR.

Whoops, I misread eltype as typeof.

I edited my previous post to have function f(x::AbstractArray{T}) where T, and now it infers to Vector{Float64} (and does not throw).

2 Likes