Type instability when broadcasting anonymous function

I found the following behavior a bit strange. When defining a function anonymously, broadcasting it introduces type instability?

julia> f = x -> x^2 # Simple function
#1 (generic function with 1 method)

julia> @code_warntype f(5) # non-broadcasting works fine
MethodInstance for (::var"#1#2")(::Int64)
  from (::var"#1#2")(x) @ Main REPL[1]:1
Arguments
  #self#::Core.Const(var"#1#2"())
  x::Int64
Body::Int64
1 ─ %1 = Main.:^::Core.Const(^)
β”‚   %2 = Core.apply_type(Base.Val, 2)::Core.Const(Val{2})
β”‚   %3 = (%2)()::Core.Const(Val{2}())
β”‚   %4 = Base.literal_pow(%1, x, %3)::Int64
└──      return %4


julia> @code_warntype f.(5) # Broadcasting it results in ::Any
MethodInstance for (::var"##dotfunction#230#3")(::Int64)
  from (::var"##dotfunction#230#3")(x1) @ Main none:0
Arguments
  #self#::Core.Const(var"##dotfunction#230#3"())
  x1::Int64
Body::Any
1 ─ %1 = Base.broadcasted(Main.f, x1)::Union{Base.Broadcast.Broadcasted{<:Base.Broadcast.BroadcastStyle, Nothing, Int64, Tuple{}}, Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}, Nothing, _A, Tuple{Int64}} where _A}
β”‚   %2 = Base.materialize(%1)::Any
└──      return %2

However, when defining the function β€œnormally” it works fine.

julia> function f2(x)
           x^2
       end
f2 (generic function with 1 method)

julia> @code_warntype f2.(5)
MethodInstance for (::var"##dotfunction#231#4")(::Int64)
  from (::var"##dotfunction#231#4")(x1) @ Main none:0
Arguments
  #self#::Core.Const(var"##dotfunction#231#4"())
  x1::Int64
Body::Int64
1 ─ %1 = Base.broadcasted(Main.f2, x1)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}, Nothing, typeof(f2), Tuple{Int64}}
β”‚   %2 = Base.materialize(%1)::Int64
└──      return %2

AFAIK I’m not doing anything weird surrounding f that should cause type instability?

I first noticed this when investigating some allocations related to a struct holding a simple function like f. When I replaced broadcasting with for-loops allocations were reduced to expected levels and @code_warntype was happy. Have I missed something about anonymous functions or is this known?

Maybe the reason is that f is not const in your anonymous function. Broadcasting is a higher order function applied on your function and in general if you include non const globals in your code you’re likely to get less than optimal inference etc

Ah thank you! In my mind f and f2 were both functions but f is, of course, just a variable holding a function.

Note, in particular, that f is a non-const global variable. If you were doing this inside a function it would be perfectly fine:

function myfunc(a)
    f = x -> x^2
    return f.(a)
end

is type-inferred for e.g. myfunc(5) or myfunc([5,6]).

The most important performance rule in Julia is to be careful with globals. Note that declaring a function with myfunc(x) = ... or function myfunc(x) ... end makes it const automatically.