8g. Type Stability with Higher-Order Functions

Introduction

Recall that functions in Julia are first-class objects, meaning that they can be manipulated as any other variable. [note] Equivalently, it's said that functions are first-class citizens. This enables us to define vectors of functions, create functions that return other functions, and more.

In particular, it enables the definition of high-order functions: functions that take another function as an argument. We've already worked with several examples of them, where anonymous functions were used as an input (e.g., map and filter).

Next, we consider the conditions under which high-order functions are type-stable.

No Specialization and Its Solutions

To illustrate the scenario, suppose we want to transform a vector x using map, and then sum its transformed elements. Additionally, assume we want to keep the transforming function generic, for which the transforming function is defined as an argument f. The implementation of this task is shown below, where we illustrate its usage through the function abs as the transformed function f.

x = rand(100)

function foo(f, x)
    y = map(f, x)
    

    sum(y)    
end

The situation is especially tricky because @code_warntype fails to detect any type-stability issues. The reason is that @code_warntype evaluates type stability under the assumption of specialization, which doesn't hold in our example: Julia bypasses specialization in the particular case where f is not explicitly called. Specifically, f in our example only enters foo as an argument of map.

To corroborate the absence of specialization, let's compare the original function with a modified version that explicitly calls f.

x = rand(100)

function foo(f, x)
    y = map(f, x)
    

    sum(y)    
end

48.447

x = rand(100)

function foo(f, x)
    y = map(f, x)
    f(1)                # irrelevant computation to force specialization

    sum(y)
end

48.447

We can observe a significant reduction in time by adding f(1), along with a decrease in memory allocations. Excessive allocations are a common symptom of type instability, as we'll demonstrate when exploring the subject.

Nonetheless, explicitly calling the function to circumvent the no-specialization issue isn't optimal, since it introduces an unnecessary computation. Fortunately, alternative solutions exist to address the problem. One approach is to type-annotate f, which provides Julia with a hint to specialize. Another solution is to wrap the function in a tuple and then call it. This ensures the identification of the function's type, as the information is inherent to the tuple's type.

For the sake of completeness, we next outline all the approaches, including the previous method with an explicit function call.

x = rand(100)


function foo(f, x)
    y = map(f, x)
    f(1)                # irrelevant computation to force specialization

    sum(y)
end

48.447

x = rand(100)


function foo(f::F, x) where F
    y = map(f, x)
    

    sum(y)
end

48.447

x = rand(100)
f_tup = (abs,)

function foo(f_tup, x)
    y = map(f_tup[1], x)
    

    sum(y)
end

48.447