Functional Programming for AI Engineers Part I: Programming Without (Completely) Losing It

By Dio (@interleave.love)
Published:

Sometimes it is really difficult to motivate thinking in a different way, to motivate approaching ideas from a different direction. Yet, it is in that journey to the same location that we may reveal paths that were difficult to discover, and even more difficult to traverse otherwise. I decided to give this talk so that you may find paths you would have otherwise.

Functional Programming for AI Engineers

Motivation

I fell in love with functional programming after struggling with systems written procedurally and object oriented for years. When things went wrong I would spend a long time having to debug and reproduce the issues that I was dealing. There was also I bunch of issues determining exactly where the problem even was. I spent an entire month improving the error handling in our applications just so that it made it easier for engineers to create robust and more easily understandable traces for our programs to fall apart.

If you are wondering why you should even be paying attention to the idea of functional programming in the first place that is understandable. Some of you may believe, 'hey, why would I need to learn about functional programming I don't write foundation models from scratch or train models from scratch.' If that is you I have to remind you that at the end of the day even if you treat AI Models like a black box they are code, and they are also code that you are relying on. If you are a vibe coder, you still have code that you will need to tweak and maintain or that you will pay someone to tweak and maintain. If you have too few engineers try to work with a system you don't understand and need to pivot you might continously end up having to start from zero without a good foundation.

For those of you who are tweaking AI models and systems like me, or possibly even rewriting systems from scratch this is going to be an easier sell. For those that are interested in why I would be tweaking AI systems in the first place, sit tight and we will get to an example use case at the end.

What is Functional Programming

It really is one thing. Programs must behave like true mathematical functions! Functions in math are more similar to a dictionary where you get a determined input for a given output. In fact, the input and output totally explain what a function is because they are all that a function is. Anything that gets in the way of that is not a real function in Math. You could break that down into two subsections of things: Don't mutate variables and don't have side effects. Which lends you to just ask what is a side-effect, side-effects are something everyone who has programmed have used. Generating a random variable, logging to the console, reading/writing to a file, these are all examples of side-effects.

What we normally call functions in programming are actually procedures. Just ways of describing subprograms with names. To illustrate what this means let's learn with a few examples.

def mutable_foo(bar: Int, baz: Int) -> Int:
    bar = bar + baz
    return bar
You can't mutate variables! Every variable must only ever be declared once. Just like a dictionary, you can't change a key in the value.
def random_foo(bar: Int) -> Int:
    baz = rand()
    return bar + baz
You can't get a random value because that would mean the same input gives you different outputs. Imagine if a dictionary gave you a different thing everytime!

def log_foo(bar: Int) -> Int:
    print(bar)
    return bar
You can't print because that is behaviour is invisible from the output!

baz = 0
def global_foo(bar: Int) -> Int
    baz = 1
    return bar
You can't mutate a global variable because that would mean the function has behaviour that mutates the state of your program! Functions in Math can't redefine what a symbol is. Definitions are given before hand.

Given these restrictions, functional programs might seem useless or slow. They get around these limitations in different ways. Here is one thing that a functional compiler can do: any function call can be replaced by the result of a previous call with the same parameters, and the language guarantees that all these rearrangements will not change the program result!

Why Functional

ML models primarily run on GPUs, and GPUs have a very different architecture to normal CPUs. It is this Architecture, and the programs that result from it that motivate writing functional programs perhaps more than any other field that exists.

If you have ever created a kernel before then you know that code for a kernel is quite restrictive compared to normal code. Especially if you want it to be performant. High perfomance code on the CPU can also take advantage, though to a lesser extent, of these restrictions. If you aren't, it helps to be familiar with GPU architecture first. GPUs have a thousands of cores compared to a CPU, which even the largest CPUs released have about 192 (AMD EPYC Instinct). That said the cores are much slower and have a more restricted set of instructions that they can run.

For example, GPUs cannot do branch prediction which makes them much worse at operating on branching logic than programs that run on kernels. There are many other differences between the two such as a difference in memory layouts and access times. We will not go into these differences here; If you want to dive into the weeds of how Modern GPU flow control etc. work this is a good resource. If there is anything to grasp from kernels it is that they excel at Parallel tasks and suffer from Sequential ones.

It turns out that parallel programming is really hard to do if you think sequentially. Functional programs bridge the gap. If there is no mutation, a powerful compiler can transform a program from appearing sequential to instead one that is parallel fairly trivially; Transforming a languages AST without these guarantees is otherwise extremely difficult or outright impossible. We won't explore why yet. If you can't be patient you can read a bit here.

Beyond the models themselves, there is also the matter of Agents. Many of us won't be writing these models from scratch and instead be mostly creating Agents. It turns out that FP is also uniquely good at this as well. That is because you can describe Agents an agents behavior in terms of a graph. And you can describe a graph in terms of a matrix. What that might mean I will leave for further investigation.

Why isn't it more popular if it has these advantages?

If you want to here a pretty decent lecture on the topic here is where you want to go.

Functional programming remains relatively uncommon, because it demands a shift in thinking that most developers are not taught to make. It asks you to treat computation not as a flow of data transformations. For many, this style feels unnatural at first. The benefits only become obvious after struggling with the hidden complexity of mutable state and uncontrolled side effects. You often have to fall into the traps of traditional paradigms before functional programming starts to make sense.

Historically, the ecosystem around functional languages has been sparse. While languages like Haskell and ML predate Python, they were rarely used in production systems and mostly used in research. Tooling was inconsistent, and documentation assumed (and to an extent some still does assume) a strong background in mathematics or theory. Even as languages like JavaScript and Python matured, functional languages were still seen as experimental or slow. The performance gap has narrowed significantly, especially with compilers like GHC or tools like JAX and languages like Futhark, but the perception of slowness remains.

There is also the challenge of abstraction. FP allows you to express pretty composable ideas, but these abstractions often seem opaque or unnecessary until you are in the weeds. Upto your heels in mutable state and trying to figure out why something is undefined or None. Many powerful concepts in functional programming are hard to explain without already understanding the problem they solve. That makes adoption harder. It is still an area of active research, and although that makes it exciting, it also means that best practices are less settled. This can make functional programming feel unstable, even when it offers clarity and performance that other paradigms struggle to match.

Review and Preview

Functional programming is a powerful paradigm whose depths have not been completely been mapped. It has a rich history, and a lot of power to make our programs more elegant; Easier to understand, and faster for much less effort on the programmers part, FP is shift that is being adopted in many of the top institutions. While it is not the default paradigm yet, that is changing.

Now that the basics of FP are understood we can talk more specifically about applications and tools that can take advantage of FP. You likely have a bunch of other questions like, what are the downsides of this programming paradigm? You might be thinking about how you could use this in your particular code base, so concrete and more complex examples to dive into are what is desired. Understanding more exactly how FP programs can be made to be high performance.