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?

1 Like

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

1 Like

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.

2 Likes