My thoughts going into this were that Haskell would be more efficient because the compiler guidance would keep the AI in check as it explored. A sort of neuro-symbolic search constraint. But the ubiquity of Python, despite it being a dynamic language, means that it’s easier for LLMs to do something akin to one shotting the code.

Most of the churn seems to be the LLM trying to lookup how to do pretty simple things in Haskell. Whereas when working with Marimo pair it gets straight to the task almost without thinking about it (about 200 tokens in vs a couple of thousand in Haskell). As a result, it’s probably “cheaper” overall to use Python for arbitrary tasks at this point in time.

That said, the churn of the trial and error of dynamic languages might flip the economics for much larger projects. That’s still an open question that requires more thorough investigation. As LLMs get “better” and more software is written in more esoteric languages the gap might close. Or, more ambitiously, the gap might close if language communities trained and evaluated open source models (ethically) that are good at navigating their specific ecosystems.

You can see how the LLM solves each task in the videos below.

This investigation will continue more systematically in subsequent posts but I think as small as this example was it’s pretty telling.

So let’s give the example a spin and see if Haskell feels clunky compared to R.

NB: Polars isn’t clunky and you should probably use that instead of Pandas but we compare to Pandas here because the original author does.

Reading the data

According to the article, we’ve been given a table of purchases from various countries, where the customer could have received a discount. We want to answer some questions about the data but we do so incrementally. The blog post showcases the ease of writing and changing business logic with dataframes.

So we start by looking at the total amount gathered from sales.

library(tidyverse)

purchases <- read_csv("purchases.csv")
purchases |> head()

With dataframe (in a cabal script) things look roughly similar give or take a few lines of code.

#!/usr/bin/env cabal
{- cabal:
  build-depends: base >= 4, dataframe
-}
module Main where

import qualified DataFrame as D

main :: IO ()
main = do
    df <- D.readCsv "./data/purchases.csv"
    print (D.take 10 df)

Pandas looks pretty similar:

import pandas as pd

purchases = pd.read_csv("purchases.csv")
purchases.head()

And finally in Frames we have:

{-# LANGUAGE TemplateHaskell #-}

module Main (main) where

import           Frames
import qualified Data.Foldable as F

tableTypes "Purchases" "purchases.csv"

loadBenchmarks :: IO (Frame Purchases)
loadBenchmarks = inCoreAoS (readTable "purchases.csv")

main :: IO ()
main = do
  ms <- loadBenchmarks
  print (take 10 (F.toList ms))

Everyone is doing well so far! We can read the CSV and show the first few rows.

Summing by amount

Now, let’s see how much we made in total sales.

purchases$amount |> sum()

R expresses the logic pretty directly. As does pandas:

purchases["amount"].sum()

For dataframe we would add the following expression:

print (D.sum (F.col @Int "amount") df)

That’s some clunk creeping in. We not only had to get the column name right, we had to remember its type. We can pull out some machinery to make it less clunky.

#!/usr/bin/env cabal
{- cabal:
  build-depends: base >= 4, dataframe
-}
{-# LANGUAGE TemplateHaskell #-}

module Main where

import qualified DataFrame as D
import qualified DataFrame.Functions as F

$(F.declareColumnsFromCsvFile "./data/purchases.csv")

main :: IO ()
main = do
    df <- D.readCsv "./data/purchases.csv"
    print (D.take 10 df)
    print (D.sum amount df)

So we traded off fragility for complexity. For this small example it doesn’t buy us much but in the next part it will.

Frames already had TemplateHaskell enabled so we just have to add:

import Lens.Micro.Extras
import Control.Foldl (fold, sum)
import Prelude hiding (sum)

-- All the old code.
    print $ fold sum (view amount <$> ms)

This is still pretty readable if you understand FP. In fact it’s a great declarative programming model that shows you what a sum is. For some problems this is great. But as we tack on complexity this will come back to bite us.

Grouping by country

So apparently they wanted the amounts by country. Now you have to go back and change the code. R still looks stunning:

purchases |>
  group_by(country) |>
  summarize(total = sum(amount))

This has a great SQL-ish API. Python is similar but starts to be a little clunky since it requires you to think about indices:

(purchases
  .groupby("country")
  .agg(total=("amount", "sum"))
  .reset_index()
)

dataframe’s TemplateHaskell setup now pays off.

print $  purchases
      |> D.groupBy [F.name country]
      |> D.aggregate ["total" .= F.sum amount]

groupBy takes a list of strings. We can recover the column’s string name with F.name. We also introduce the expression DSL that we use to calculate the sum (instead of the top-level D.sum function). There’s a little complexity but now we get a SQL-ish API that doesn’t leak that much implementation detail.

Frames starts to bring out the artillery.

{-# OPTIONS_GHC -Wall            #-}

{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ExplicitForAll      #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}

module Main (main) where

import           Frames
import qualified Data.Foldable as F
import qualified Data.List as L
import           Lens.Micro.Extras

import           Control.Foldl (fold, sum)
import qualified Control.Foldl as Foldl
import           Prelude hiding (sum)

import qualified Frames.MapReduce as FMR
import qualified Frames.Folds as FF

tableTypes "Purchases" "purchases.csv"

loadBenchmarks :: IO (Frame Purchases)
loadBenchmarks = inCoreAoS (readTable "purchases.csv")

unpack :: FMR.Unpack Purchases Purchases
unpack = FMR.unpackFilterOnField @Country (const True)

assign :: FMR.Assign (Record '[Country]) Purchases (Record '[Amount, Discount])
assign = FMR.splitOnKeys @'[Country]

reduce :: FMR.Reduce (Record '[Country])
                     (Record '[Amount, Discount])
                     (Frame Purchases)
reduce = FMR.foldAndAddKey $ (FF.foldAllConstrained @Num @'[Amount, Discount]) sum

mrFold :: FMR.Fold Purchases (Frame Purchases)
mrFold = FMR.concatFold $ FMR.mapReduceFold unpack assign reduce

rhead :: Show a => Frame a -> IO ()
rhead  = \ms -> mapM_ print (((take 6) . F.toList) ms)

main :: IO ()
main = do
  ms <- loadBenchmarks
  rhead ms
  print $ fold sum (view amount <$> ms)
  let result = FMR.fold mrFold ms
  putStrLn $ (L.intercalate "\n" $ fmap show $ fold Foldl.list result)

At the points at which computation happens it’s pretty declarative but simple logic is now overwhelmed by boilerplate. Refactoring to answer a pretty easy follow up question cost us a lot in code but we get compile-time safety. At this stage, how much complexity you’re willing to stomach is a matter of person taste and familiarity.

Deducting the discount

Frames is now clunky so we’ll stop reporting on it. We now want to get the sum of not just the amount but the amount minus the discount. Let’s see how the other implementations handle the discount.

# R
purchases |> 
  group_by(country) |> 
  summarize(total = sum(amount - discount))

The R change is pretty easy. Pandas gets a little weird:

(purchases
  .groupby("country")
  .apply(lambda df: (df["amount"] - df["discount"]).sum())
  .reset_index()
  .rename(columns={0: "total"})
)

We now start leaking a few more implementation details. Column renaming and the difference between simple sums and expression sums.

Haskell still tracks R:

print $  purchases
      |> D.groupBy [F.name country]
      |> D.aggregate ["total" .= F.sum (amount - discount)]

Removing outliers

Let’s now remove any rows whose amount is greater than the global median * 10.

purchases |>
  filter(amount <= median(amount) * 10) |>
  group_by(country) |> 
  summarize(total = sum(amount - discount))

Pandas introduces a query function:

(purchases
  .query("amount <= amount.median() * 10")
  .groupby("country")
  .apply(lambda df: (df["amount"] - df["discount"]).sum())
  .reset_index()
  .rename(columns={0: "total"})
)

dataframe looks fairly similar too:

print $  purchases
      |> D.filterWhere (F.toDouble amount .<= F.median amount * 10)
      |> D.groupBy [F.name country]
      |> D.aggregate ["total" .= F.sum (amount - discount)]

amount is an Int and median is a double so we pay a small conversion cost but that’s okay.

Things look good on all fronts. The refactor wasn’t hard.

Using the median within each country

R just drops the comparison to after the groupBy.

purchases |>
  group_by(country) |>
  filter(amount <= median(amount) * 10) |>
  summarize(total = sum(amount - discount))

Pandas struggles to handle this logic succintly:

(purchases
  .groupby("country")
  .apply(lambda df: df[df["amount"] <= df["amount"].median() * 10])
  .reset_index(drop=True)
  .groupby("country")
  .apply(lambda df: (df["amount"] - df["discount"]).sum())
  .reset_index()
  .rename(columns={0: "total"})
)

The diff now includes a lot of extra noise and takes you away fro the core change.

Haskell keeps the diff minimal:

print $  purchases
      |> D.filterWhere (F.toDouble amount .<= F.over [F.name country] (F.median amount * 10))
      |> D.groupBy [F.name country]
      |> D.aggregate ["total" .= F.sum (amount - discount)]

We use the over expression which does an intra-expression grouping. All the setup we did is now paying off. Our diffs are small and focused. Arguably, the R diff (changing the order of the groupby/filter) makes things a little more confusing because your reasoning is non-local. You need to read the previous line to understand what

What about type safety?

We lost some type safety by making the API less clunky. For exmaple, this gives us a runtime failure.

print $  purchases
      |> D.exclude [F.name amount]
      |> D.filterWhere (F.toDouble amount .<= F.over [F.name country] (F.median amount * 10))
      |> D.groupBy [F.name country]
      |> D.aggregate ["total" .= F.sum (amount - discount)]

Filter will try and refer to an amount that doesn’t exist then throw an exception. Frames would never let us do something like this. Can we get type safety back?

Yes we can.

We just have to swap out some imports and do a little more annotating.

#!/usr/bin/env cabal
{- cabal:
  build-depends: base >= 4, dataframe, text
-}
{-# LANGUAGE DataKinds #-}

module Main where

import qualified DataFrame as D
import qualified DataFrame.Typed as DT

import Data.Text (Text)

main :: IO ()
main = do
    df <- D.readCsv "./data/purchases.csv"
    let tdf = DT.unsafeFreeze @'[ DT.Column "country" Text
                                , DT.Column "amount" Int
                                , DT.Column "discount" Int] df

    print $  tdf
          |> DT.filterWhere (DT.toDouble (DT.col @"amount") .<= DT.over @'["country"] (DT.median (DT.col @"amount" * 10))
          |> DT.groupBy @'["country"]
          |> DT.aggregate (DT.agg @"total" (DT.sum (DT.col @"amount" - DT.col @"discount") DT.aggNil))

We could use template Haskell to generate the type we freeze the dataframe into.

-- this is the equivalent of writing:
-- type Purchases = '[ DT.Column "country" Text
--                   , DT.Column "amount" Int
--                   , DT.Column "discount" Int]
$(DT.deriveSchemaFromCsvFile "Purchases" "./data/purchases.csv")


-- Skip other code
    let tdf = DT.unsafeFreeze @Purchases df

We’re only a little clunkier than before. We could in principle add more TemplateHaskell to get the column references similar to what we did to remove F.col before. Leaving us with:

-- Some code above.
    print $  tdf
          |> DT.filterWhere (DT.toDouble amount .<= DT.over @'["country"] (DT.median (amount * 10)))
          |> DT.groupBy @'["country"]
          |> DT.aggregate (DT.agg @"total" (DT.sum (amount - discount) DT.aggNil))

Up shot

Safety and usability are usually at tension. But you can use a compiled language (in this case Haskell) as ergonomically as you would a dynamic language if you design the right abstractions and keep them centered on user tasks.

A notebook with a visual TLDR is available as an export of the marimo notebook. You can import the raw notebook into molab to try it out yourself.

The first idea comes from a problem I kept running into at work.

Boosting

I spent most of my career working in fraud detection. XGBoost (or similar techniques) dominate in this space. XGBoost approximates a function by training small decision trees to incrementally chip away at a residual. Each tree doesn’t do a great job on its own, but they pick up each other’s slack by collectively contributing towards the whole picture. There’s an old parable about blind men touching different parts of an elephant. One feels the trunk and says it’s a snake, another feels the leg and says it’s a tree, a third feels the side and says it’s a wall. Boosting works like this. Each tree touches a different part of the function and gets a partial answer. Their sum reaches the right conclusion. (There is a better treatment of boosting in general here.)

But similar to the parable, the trees end up with a lot of detail that’s extraneous to the underlying function. Each individual tree tells you little about the final prediction. The ensemble is accurate but a little opaque. You can use game-theoretic techniques like SHAP to determine what each feature contributes to the output, and that gives you a coarse view of what the model is doing. But the most interpretable form a model can come in is a single tree. That’s just nested if-statements. You wouldn’t have to probe the model in creative ways. You could follow the lines of implication and see exactly why a decision was made.

I like this human-centric way of building models. The model is something to be understood, changed, questioned, inspected. There are practical benefits too. If you can translate a model into a single decision tree, you can translate it to SQL, to Go, to whatever your production environment uses. You’re not locked into serving the model in the framework it was trained in.

So how do you go from an ensemble of hundreds of trees to a single tree?

Turns out that problem is NP-hard. A compact tree reconstructed from an ensemble is called a born-again tree. The technique uses dynamic programming to regenerate a single tree that faithfully represents the ensemble. But it’s computationally expensive, and it might give you a tree that tries to fit the noise of the ensemble. After all, learning is a messy process.

You could also distill the ensemble: train a new tree not against the original training data but against the ensemble’s predictions. The student learns from the teacher. But even learning an optimal decision tree is NP-hard, so you fall back to greedy algorithms like CART. The process is lossy. You’re not sure what detail it misses. It captures the main essence, but you’ve traded some accuracy for interpretability without a clear guarantee of what you kept and what you lost.

So we’re kind of stuck. The ensemble is accurate but unreadable. Collapsing it exactly is intractable. Distilling it is lossy and hard to reason about.

Let’s hold on to this thought and take a detour.

Symbolic regression

Symbolic regression is a different approach to learning from data entirely. Instead of fitting parameters of a fixed model (like linear regression fitting coefficients), it searches over the space of mathematical expressions to discover the structure itself. The output isn’t a weight matrix or a tree. It’s an equation: $y = 2.5\,x^2 + 0.3\,x$, or $F = G\,m_1\,m_2 / r^2$.

Most symbolic regression methods use genetic programming. You generate a population of candidate expressions, evaluate them against the data, run tournaments where the best survive, crossover and mutate the survivors, and repeat. You are effectively trying to one-shot the learning problem: find the right expression structure and the right constants in a single search.

This works on small-ish datasets and is good at recovering the equation that describes the underlying data. Symbolic regression has been popular in the sciences because it promises testable hypotheses rather than a black box that produces just answers. It’s the process Kepler went through trying to find laws of planetary motion, except automated. In a podcast with Terrence Tao, Dwarkesh Patel described Kepler’s process as being like a “high temperature LLM”. The analogy captures the trial-and-error aspect but misses the heart of it. Symbolic methods are greatly improved by a combination of reasoning and search. PySR, the most popular tool for this, runs multi-population evolutionary search and recovers exact physical laws from noisy data in seconds.

There’s also recent work by Fabricio Olivetti de França on eggp, a symbolic regression system that uses e-graphs to prune the search space. We’re already a few digressions deep, but the short version is that e-graphs are a data structure that stores many equivalent forms of an expression simultaneously, and after applying algebraic rewrite rules, gives you back the simplest one. Think of it as a compiler optimization pass for mathematical expressions.

The convergence

So XGBoost has this property where an ensemble of trees is accurate but hard to collapse into something readable. And symbolic regression typically tries to one-shot learn a single equation from data. What if we combined the two approaches?

The idea is: use boosting to build an ensemble of small symbolic expressions (not trees, but tiny equations like 0.3*x^2 or sin(x0)), then collapse the ensemble algebraically into a single clean formula. Accept that learning is messy, but have techniques to consolidate it later.

This is similar to how studying works. You take a bunch of scattered notes during a lecture, full of redundancies and partial thoughts, and later consolidate them into a neater structure. Traditional symbolic regression tries to make the lecture notes and the final summary the same document. Boosting lets them be different. You scribble fast and clean up later.

It also reminds me of the wake-sleep algorithm from the DreamCoder paper on program synthesis. DreamCoder alternates between a “wake” phase (solve problems using your current library) and a “sleep” phase (compress solutions into reusable abstractions). The wake phase is messy and exploratory. The sleep phase consolidates. Symbolic boosting has the same rhythm: grow the ensemble (wake), simplify it (sleep), grow again, simplify again.

How it works

Concretely, symbolic boosting works like this. Initialize the model as the mean of the target. The residuals are everything the mean doesn’t explain. Then, each round:

1. Generate 50 random expression trees, each with at most 7 nodes. The grammar includes +, -, *, and squaring.
2. Fit each candidate’s learnable parameters against the current residuals using L-BFGS-B (30 iterations, 2 random restarts).
3. Score each candidate by how much it reduces the residual variance.
4. Pick the winner. Shrink it by a learning rate (default 0.1) and add it to the ensemble.
5. Update residuals and go again.

After 30 rounds, the ensemble might look like:

0.83*x*y + (-0.90)*x*y + 1.18*y*(x+y) - 1.01*(y-0.51)^2 - 0.42*y - 0.76

That’s the scribbled lecture notes. Redundant terms, partial cancellations, coefficients scattered everywhere. The true answer is x*y, buried under noise.

Now for the consolidation.

Simplification: polynomial normal form

If you ignore transcendental functions (sin, cos, exp, log), the boosting ensemble is a polynomial. And polynomials have a canonical form where identical monomials are merged automatically.

Take the messy ensemble above:

Expand every product and collect like terms:

Result: $xy + 0.77\,y - 1.02$. Three terms instead of seventeen.

The implementation is around 25 lines of Python using SymPy:

TRANSCENDENTALS = (sympy.sin, sympy.cos, sympy.exp, sympy.log)

def poly_nf(expr, threshold=1e-9):
    atom_map, ctr = {}, [0]
    def replace_trans(e):
        if e.func in TRANSCENDENTALS:
            inner = poly_nf(e.args[0], threshold)
            ne = e.func(inner); k = str(ne)
            if k not in atom_map:
                s = sympy.Symbol(f"atom{ctr[0]}"); ctr[0] += 1
                atom_map[k] = (s, ne)
            return atom_map[k][0]
        return (e.func(*[replace_trans(a) for a in e.args])
                if e.args else e)
    replaced = replace_trans(expr)
    expanded = sympy.expand(replaced)
    terms = sympy.Add.make_args(expanded)
    if len(terms) > 200:
        return expr
    kept = [t for t in terms
            if abs(float(t.as_coeff_Mul()[0])) > threshold]
    result = sympy.Add(*kept) if kept else sympy.Float(0)
    for key, (s, orig) in atom_map.items():
        result = result.subs(s, orig)
    return result

First, replace every transcendental subexpression (sin(...), exp(...), etc.) with a fresh placeholder symbol. This draws a boundary: everything polynomial gets expanded and merged, everything transcendental stays opaque. Then sympy.expand distributes all products over sums. Like terms merge automatically. Drop anything with a near-zero coefficient.

The transcendental boundary is important. Without it, sympy.expand would try to expand through sin(x + y) and fail. With it, sin(x + y) becomes an atom that participates in polynomial arithmetic the same way a variable does. After simplification, the placeholders are substituted back.

The blowup guard (bail out if expansion produces more than 200 terms) prevents polynomial multiplication from cascading. Two 15-term polynomials multiplied together produce up to 225 monomials. Without the guard, nested products could blow up the representation.

Simplification: magnitude pruning

Polynomial normal form is algebraic simplification. Magnitude pruning is statistical simplification.

After expansion, you might have a monomial like 0.003 * cos(y) alongside a dominant 98.7 * x * y. Magnitude pruning evaluates each monomial on the training data, measures its peak absolute contribution, and drops anything contributing less than 3% of the dominant term. The cos(y) correction is noise relative to the signal. Drop it.

def mag_prune(expr, X_data, threshold=0.03):
    expanded = sympy.expand(expr)
    terms = list(sympy.Add.make_args(expanded))
    if len(terms) <= 1:
        return expr
    fsyms = sorted(expr.free_symbols, key=str)
    peaks = []
    for t in terms:
        try:
            fn = sympy.lambdify(fsyms, t, modules=["numpy"])
            vals = fn(*[X_data[:, int(str(s)[1:])] for s in fsyms])
            peaks.append(float(np.nanmax(np.abs(
                np.nan_to_num(vals, nan=0.0)))))
        except Exception:
            peaks.append(float("inf"))
    mx = max(peaks) if peaks else 1.0
    if mx == 0:
        return expr
    kept = [t for t, p in zip(terms, peaks) if p / mx >= threshold]
    return sympy.Add(*kept) if kept else expr

Simplification: e-graph cleanup

Polynomial normal form and magnitude pruning handle polynomial redundancy and statistical insignificance. But there’s a third kind of redundancy: algebraic identities involving transcendental functions.

If the ensemble contains log(exp(x)), polynomial normal form treats it as an opaque atom and leaves it alone. It doesn’t know that log(exp(x)) = x. Similarly, sqrt(x)^2 = x, x + 0 = x, and x * 1 = x are identities that polynomial algebra can’t simplify.

This is where e-graphs come back. The notebook uses egglog to apply a conservative set of rewrite rules:

erules = [
    rewrite(a + MathE(0.0)).to(a),        # x + 0 → x
    rewrite(a * MathE(1.0)).to(a),         # x * 1 → x
    rewrite(a * MathE(0.0)).to(MathE(0.0)),# x * 0 → 0
    rewrite(-(-a)).to(a),                  # double negation
    rewrite(Log(Exp(a))).to(a),            # log(exp(x)) → x
    rewrite(Exp(Log(a))).to(a),            # exp(log(x)) → x
    rewrite(Sqrt(a) ** MathE(2.0)).to(a),  # sqrt(x)^2 → x
    rewrite(a * a).to(a ** MathE(2.0)),    # x*x → x^2
]

Insert the expression, run saturation for 30 iterations, extract the smallest equivalent.

The rules are deliberately conservative. Distributivity (a*(b+c) = a*b + a*c) is not included because it causes e-graphs to explode. Polynomial normal form already handles distributivity structurally, so the e-graph only needs to handle what polynomials can’t.

The sawtooth

The boosting loop doesn’t simplify every round. It simplifies every 5 rounds. The ensemble grows messy for 5 rounds, then collapses when the simplification pipeline fires. Then it grows messy again, then collapses again.

This creates a sawtooth pattern in expression complexity. The notebook visualizes this: the top plot shows complexity spiking up and crashing down at each collapse event, the bottom plot shows R-squared climbing steadily through the collapses. The collapses don’t hurt accuracy because simplification doesn’t change semantics. Merging 0.83*x*y + (-0.90)*x*y into -0.07*x*y produces the exact same predictions on every input. The expression gets smaller, not different.

The rhythm is: grow, simplify, grow, simplify. Scribble, consolidate, scribble, consolidate.

Where it works and where it doesn’t

The notebook includes four target functions to try boosting against.

2.5*x^2 + 0.3*x (simple quadratic). A polynomial. Boosting nails it. The ensemble discovers x^2 and x terms early, and polynomial normal form collapses them into the exact answer.

sin(x1) + cos(x2) (trigonometric). Also additive. Each trig function is a plausible weak learner on its own. Boosting finds the first term, the residuals look like cos(x2), and it finds the second.

3.0*x1^2 + 0.5/x2 (physics-inspired). The polynomial part is easy. The 0.5/x2 is a single division. If the random generator produces something like param0 / x1, parameter fitting finds the right coefficient.

x1 * exp(-x2^2) + 2.3 (complex). This is the hard one. exp(-x2^2) is a deeply nested composition. No single 7-node weak learner can express x1 * exp(-x2^2). The ensemble will approximate it with polynomial corrections, but it’s unlikely to recover the exact structure.

This is the fundamental limitation. Some mathematical structures are not additive. exp(-x^2) has to be built as a single composite expression, not assembled from a sum of corrections. Boosting will converge to a good numerical approximation, but the formula it produces will be a polynomial approximation of the true equation, not the equation itself.

PySR’s evolutionary approach doesn’t have this bias. It searches over full expression trees, so it can discover exp(-x^2) as a single unit. The tradeoff: PySR needs more time and is more stochastic. Boosting is faster and more predictable, but structurally limited to additive decompositions.

The sweet spot is problems where the underlying equation is a sum of a few terms, each of moderate complexity. Coulomb’s law (q1*q2 / r^2), kinetic energy (0.5*m*v^2), spring force (k*x). Additive or multiplicative structure that the polynomial representation can collapse.

Try it yourself

The three tools that make this possible: PySR for evolutionary symbolic regression, SymPy for symbolic algebra and the polynomial normal form, and egglog for e-graph equality saturation.

Marimo notebooks are great for exploring this sort of approach since they make visualizing the outputs easy. We can leverage human intelligence to further audit and simplify models.

Why this is exciting

Symbolic regression has always been hamstrung by scale. Finding ways to make training loops smaller and more efficient removes an important hurdle to its adoption. Finding different ways to create symbolic models that search for or approximate exact equations might be a chapter in the AI story where we actually think about how we can collaborate with machines to discover new things.

I loaded it up in GHCi and ran the untyped version first:

dataframe> df <- D.readCsv "./data/creditcard.csv"
dataframe> :set +s
dataframe> :declareColumns df
time :: Expr Double
v1 :: Expr Double
v2 :: Expr Double
...
_class_ :: Expr Int
dataframe> df |> D.groupBy [F.name _class_]
         |> D.aggregate ["rand" .= F.mean v1 - F.mean v2]
         |> D.select ["rand"]
---------------------
        rand
---------------------
       Double
---------------------
-8.395726543461365
1.4528594901372989e-2

(0.92 secs, 4,287,695,984 bytes)

But suppose you had a runtime error in the last part of the pipeline?

dataframe> df |> D.groupBy [F.name _class_]
         |> D.aggregate ["rand" .= F.mean v1 - F.mean v2]
         |> D.select ["ran"]

*** Exception: 

[ERROR] Column not found: ran for operation select
	Did you mean rand?

Now you’ve waited the duration of the entire execution of the pipeline (allocations and all) to discover that there is a misnamed column. That’s pretty wasteful.

So I ported the example to the typed dataframe and API and ran it. First the misnamed version:

The example did show that when you misname the column it fails at compile time.

dataframe> :set -XDataKinds
dataframe> import qualified DataFrame.Typed as DT
dataframe> _ = (); DT.deriveSchema "CreditCard" df. -- Template Haskell trick for GHCi
dataframe> tdf = either (error . show) id (DataFrame.Typed.freezeWithError @CreditCard df)
dataframe> tdf |> DT.groupBy @'["Class"]
               |> DT.aggregate (DT.agg @"rand" (DT.mean (DT.col @"V1") - DT.mean (DT.col @"V2")) DT.aggNil)
               |> DT.select @'["ran"]
<interactive>:39:127: error: [GHC-64725]
    • Column 'ran' not found in schema
    • In the second argument of ‘(|>)’, namely ‘DT.select @'["ran"]’
      In the expression:
        tdf |> DT.groupBy @'["Class"]
          |>
            DT.aggregate
              (agg
                 @"rand" (DT.mean (DT.col @"V1") - DT.mean (DT.col @"V2")) aggNil)
          |> DT.select @'["ran"]
(0.00 secs,)

And then the happy path:

dataframe> tdf |> DT.groupBy @'["Class"]
               |> DT.aggregate (DT.agg @"rand" (DT.mean (DT.col @"V1") - DT.mean (DT.col @"V2")) DT.aggNil)
               |> DT.select @'["rand"]

Nothing. The cursor just sat there. I waited a minute. Two minutes. Still nothing. My friend was polite about it and started conjecturing about what could be wrong. But the happy path just didn’t finish and eventually OOM’d.

Chasing ghosts

I spent a while convinced the problem was runtime. Both APIs generate the same expression tree under the hood. Both call the same compiled aggregation function. The typed wrappers are just newtypes. So I went looking for specialization failures, thinking GHCi wasn’t generating efficient code for the polymorphic hot loops.

I added SPECIALIZE pragmas to foldLinearGroups, mapColumn, zipWithColumns, and every aggregation function in Functions.hs. I made the typed combinators delegate to the untyped compiled versions instead of defining their own closures. I specialized for Double, Float, Int, Int8, Int16, Int32, Int64. The whole works.

No change. Still hung.

At this point I wrote both pipelines into a compiled executable with -O2 and profiling enabled. Both finished in 13 milliseconds. Identical. The runtime code was not the problem. Whatever was happening was specific to GHCi.

Then I did the thing I should have done from the start. I put Debug.Trace calls at every stage of the pipeline:

aggregate tagg (TGD gdf) =
    trace "[typed.aggregate] START" $
    unsafeFreeze (DA.aggregate (taggToNamedExprs tagg) gdf)

aggregate aggs gdf@(Grouped df ...) =
    trace "[aggregate] START" $
    ...

eval ctx expr@(Binary ...) =
    trace "[eval] Binary" $
    ...

None of them fired. The program never reached runtime.

I tried a few more things after that. I thought maybe it was laziness, that evaluation was being deferred and the traces weren’t forced. I added bang patterns, seq calls, explicit evaluate in IO. Nothing changed. The traces weren’t printing because the code wasn’t running.

That’s when I understood: the hang was happening during GHCi’s type-checking of the expression at the prompt.

The confirmation

To be sure, I ran the same typed pipeline against the housing dataset. 10 columns instead of 31.

[typed.aggregate] START
[typed.aggregate] exprs: [("rand",(sub (mean (col @Double "median_income"))
                                       (mean (col @Double "housing_median_age"))))]
[aggregate] START
[eval] Binary
[eval] MergeAgg FAST PATH col="median_income"
[eval] MergeAgg FAST PATH col="housing_median_age"
[eval] Binary: liftValue2
[aggregate] expr DONE
(0.55 secs, 410,398,240 bytes)

Everything worked. All the traces fired. Fast path, correct result. So the runtime code was fine. The problem scaled with the number of columns in the schema. 10 columns: instant. 31 columns: heat death of the universe.

The actual bug

The typed API uses type families to compute result schemas. When you call groupBy @'["Class"] on a 31-column schema, the type checker needs to figure out which columns are grouping keys. That computation was done by a type family called GroupKeyColumns which itself uses a type-level If statement to check for column existence:

type family GroupKeyColumns (keys :: [Symbol]) (cols :: [Type]) :: [Type] where
    GroupKeyColumns keys '[] = '[]
    GroupKeyColumns keys (Column n a ': rest) =
        If
            (IsElem n keys)
            (Column n a ': GroupKeyColumns keys rest)
            (GroupKeyColumns keys rest)

It’s easy to think of type-level programming as some kind of compiler magic that either works or doesn’t. But type families are a sort of programming language that must be reasoned about with the same tools we use to reason about regular programs. They take arguments, do recursion, and return results.

So let’s trace through what GroupKeyColumns '["Class"] does on a 3-column schema, step by step. This is the If version:

type family If (cond :: Bool) (t :: k) (f :: k) :: k where
    If 'True  t f = t
    If 'False t f = f

GroupKeyColumns '["Class"] '[Column "V1" Double, Column "V2" Double, Column "Class" Int]

-- Step 1: n = "V1". IsElem "V1" '["Class"] = 'False.
-- GHC needs to evaluate If 'False <then> <else>.
-- But first it normalizes both arguments:
--   then = Column "V1" Double ': GroupKeyColumns '["Class"] '[Column "V2" Double, Column "Class" Int]
--   else = GroupKeyColumns '["Class"] '[Column "V2" Double, Column "Class" Int]
-- Both branches trigger the recursion. Two copies of the same work.

-- Step 2 (from the "then" branch): n = "V2". Same story. Two more copies.
-- Step 2 (from the "else" branch): n = "V2". Same story. Two more copies.
-- We now have 4 pending evaluations of GroupKeyColumns on '[Column "Class" Int].

-- Step 3: each of those 4 evaluates n = "Class". IsElem = 'True.
-- Each produces two more recursive calls on '[].
-- 8 evaluations of GroupKeyColumns '["Class"] '[] = '[].

Three columns, eight leaf evaluations. The tree of work doubles at every level. For n columns, GHC does 2^n reductions. Surprisingly, type-level if statement aren’t lazy as I think I expected them to be.

With 10 columns (the housing dataset I tested against), that’s 1,024 reductions. Imperceptible. With 31 columns (the Kaggle credit card dataset I demoed with), that’s 2,147,483,648 reductions. That’s why it hung.

The same analysis applies to any recursive type family that uses If with recursive calls in both branches. It’s the type-level equivalent of writing a naive Fibonacci function without memoization. The recursion tree fans out because GHC doesn’t share work between the two If arguments.

The fix

I replaced If with a helper type family that pattern-matches directly on the boolean:

type family GroupKeyColumns (keys :: [Symbol]) (cols :: [Type]) :: [Type] where
    GroupKeyColumns keys '[] = '[]
    GroupKeyColumns keys (Column n a ': rest) =
        GroupKeyColumnsHelper (IsElem n keys) n a keys rest

type family GroupKeyColumnsHelper (found :: Bool) (n :: Symbol) (a :: Type)
    (keys :: [Symbol]) (rest :: [Type]) :: [Type] where
    GroupKeyColumnsHelper 'True  n a keys rest =
        Column n a ': GroupKeyColumns keys rest
    GroupKeyColumnsHelper 'False n a keys rest =
        GroupKeyColumns keys rest

Now only the matching branch is ever constructed. The recursion is linear.

After the change the typed and untyped version run at roughly the same speed.

Why the helper version works

Let’s trace the helper version on the same 3-column schema:

GroupKeyColumns '["Class"] '[Column "V1" Double, Column "V2" Double, Column "Class" Int]

-- Step 1: n = "V1". IsElem "V1" '["Class"] = 'False.
-- GHC evaluates: GroupKeyColumnsHelper 'False "V1" Double '["Class"] '[Column "V2" Double, Column "Class" Int]
-- Matches the 'False equation. Result: GroupKeyColumns '["Class"] '[Column "V2" Double, Column "Class" Int]
-- Only one recursive call. The 'True branch was never constructed.

-- Step 2: n = "V2". IsElem "V2" '["Class"] = 'False.
-- Same thing. One recursive call.

-- Step 3: n = "Class". IsElem "Class" '["Class"] = 'True.
-- Matches the 'True equation. Result: Column "Class" Int ': GroupKeyColumns '["Class"] '[]
-- One recursive call.

-- Step 4: GroupKeyColumns '["Class"] '[] = '[]

Four steps total. Linear. Each step does constant work because only the matching equation is ever expanded. The other equation doesn’t exist in the reduction trace at all.

This is a known footgun in type-level Haskell. But If is the natural thing to reach for when you want a conditional, and it works fine on small inputs. I tested everything against the 10-column housing dataset and it was fast. I didn’t think to test against 31 columns until I was demoing it live.

The rule of thumb

If you have a recursive type family and you’re using If to select between branches, you have an exponential blowup waiting for a wide enough input. My fix was to introduce a helper type family with two equations and let GHC’s pattern matching do the branching instead of If.

-- Exponential: both branches are always reduced
type family Foo (xs :: [k]) :: [k] where
    Foo '[] = '[]
    Foo (x ': xs) = If (SomePredicate x) (x ': Foo xs) (Foo xs)

-- Linear: only the matching branch is constructed
type family Foo (xs :: [k]) :: [k] where
    Foo '[] = '[]
    Foo (x ': xs) = FooHelper (SomePredicate x) x xs

type family FooHelper (b :: Bool) (x :: k) (xs :: [k]) :: [k] where
    FooHelper 'True  x xs = x ': Foo xs
    FooHelper 'False x xs = Foo xs

It’s more verbose but it’s O(n) instead of O(2^n). For 10 columns you won’t notice. For 31 columns you’ll be sitting there wondering why your program stopped responding while your friend (who is convinced that Haskell is an impractical academic fascination) watches.

I ran into this question while building my own dataframe library. I needed to decide which operations were truly fundamental and which were just surface-level variations. That search led me to Petersohn et al.’s Towards Scalable Dataframe Systems. They’d built Modin, a drop-in replacement for pandas, and needed to understand the actual structure underneath the API. They analyzed 1 million Jupyter notebooks, cataloged how people use pandas, and proposed a dataframe algebra: a formal set of about 15 operators that can express what all 200+ pandas operations do.

That algebra was a huge compression, but I kept wondering whether there was a level below it. Whether there was a smaller set of truly primitive operations that the 15 are built from. If those exist, they would be a real foundation: operations small enough to be obviously correct, expressive enough to build everything else.

Petersohn’s dataframe algebra

It’s worth spending some time on what Petersohn et al. actually did, because it frames everything that follows.

They started by defining what a dataframe is. Surprisingly, nobody had done this formally before. Their Definition 4.1 says a dataframe is a tuple (A, R, C, D): an array of data A, row labels R, column labels C, and a vector of column domains D. This is more precise than “a table” because it captures things that make dataframes different from relational tables. Rows and columns are both ordered, both labeled, and treated symmetrically. You can transpose a dataframe. You can promote data values into column labels. These aren’t things you can do with a SQL table.

Then they identified the operators. Here’s their Table 1, condensed:

The “Origin” column matters. The first nine operators come from relational algebra and have direct analogs in SQL. WINDOW comes from SQL extensions. The last four (TRANSPOSE, MAP, TOLABELS, FROMLABELS) are unique to dataframes. They exist because dataframes treat rows and columns symmetrically and allow data to move between values and metadata. Relational databases can’t do that.

Petersohn showed that over 85% of the pandas API can be rewritten as compositions of these operators. Operations like fillna, isnull, str.upper, and cummax are all special cases of MAP. Operations like sort_values, set_index, reset_index, merge, groupby, and pivot all map one-to-one onto operators in the algebra. That’s a huge compression: 200+ ad hoc methods become 15 composable primitives.

But I kept looking at the relational operators in that table (PROJECTION, RENAME, GROUPBY, JOIN) and thinking: these feel related. They all change the schema of the dataframe. Is there a deeper relationship?

Shapes of schema change

I kept staring at Petersohn’s table, and a pattern emerged. Some operators change the schema, meaning which columns exist and what types they have. Others leave the schema alone and only affect the rows. And if you focus on the schema-changing ones, they fall into three groups.

Restructuring. You rearrange, subset, or relabel columns. The data stays the same; only the shape changes. In SQL terms: SELECT name, salary FROM employees produces a two-column result from a three-column table. In the dataframe library:

select ["name", "salary"] df
-- 3-column schema → 2-column schema

rename "salary" "pay" df
-- Column name changes, data untouched

exclude ["department"] df
-- Drop a column

This covers Petersohn’s PROJECTION and RENAME. The output schema is a function of the input schema. You can compute it without looking at any data.

Merging. You collapse rows that share a key, either into a summary or a collection. In SQL: SELECT department, AVG(salary) FROM employees GROUP BY department. In the library:

aggregate
    [ mean (col @Double "salary") `as` "avg_salary"
    , count (col @Text "name") `as` "headcount"
    ]
    (groupBy ["department"] df)
-- Schema: name, department, salary
--       → department, avg_salary, headcount

-- Or keep all values without reducing:
aggregate
    [ collect (col @Double "salary") `as` "all_salaries" ]
    (groupBy ["department"] df)
-- Each department gets a list of all its salaries

Multiple rows map to the same key and get combined. This covers Petersohn’s GROUPBY and UNION.

Pairing. You find rows in two tables that agree on a shared key and stitch them into a wider row. In SQL: SELECT * FROM employees INNER JOIN departments USING (department). In the library:

innerJoin ["department"] employees departments
-- Schema: (name, department, salary) + (department, budget)
--       → name, department, salary, budget

Shared keys appear once; unique columns from each side are concatenated. Left and outer joins are the same idea but with nullable columns where matches might be missing. This covers Petersohn’s CROSS PRODUCT / JOIN.

What doesn’t fit. Two relational operators resist this grouping. In SQL: SELECT * FROM employees EXCEPT SELECT * FROM contractors returns rows in one table but not the other. And SELECT DISTINCT * FROM employees collapses duplicate rows. In the library:

distinct df
-- Same schema, fewer rows: removes duplicates

-- DIFFERENCE is not yet implemented, but the idea is:
-- difference employees contractors
-- Same schema, fewer rows: removes rows present in both

DIFFERENCE and DROP DUPLICATES both change which rows are present, but they don’t restructure columns, collapse by key, or pair across tables. They feel set-theoretic: one computes a complement, the other computes an image. For now I’ll set them aside, but they’ll reappear when the categorical picture gets sharper.

So five of Petersohn’s relational operators (PROJECTION, RENAME, GROUPBY, UNION, JOIN) map cleanly onto three patterns: restructuring, merging, pairing. The schema-preserving operators (SELECTION, SORT, WINDOW) are orthogonal. They change which rows you see or in what order, but not the column structure. And the dataframe-specific operators (TRANSPOSE, MAP, TOLABELS, FROMLABELS) live outside the relational model entirely.

I had three patterns and two outliers. What I didn’t have was a reason why it should be these three patterns. Are restructuring, merging, and pairing truly fundamental, or did I just happen to group things this way? Where do DIFFERENCE and DROP DUPLICATES belong? Is there a theory that predicts all of this?

That’s the question my mentor Sam Stites pointed me toward when he suggested I read Fong and Spivak’s Seven Sketches in Compositionality. The answer turns out to come from category theory. The version of category theory in Chapter 3 of that book is concrete and database-flavored. You don’t need to know abstract math. You just need to think carefully about what a schema is and what it means to change one.

Why three migration functors

Let’s start with a concrete example. Imagine you have two schemas:

Employees: name, department, salary
Departments: department, budget

There’s a natural relationship between them. The department column in Employees references the department column in Departments. That’s a foreign key. It’s a mapping from one schema to the other: every employee row points at a department row.

Now, what can you do with that mapping? Fong and Spivak’s Chapter 3 identifies three fundamental operations.

You can restrict data to fit a different schema. The schema {name, salary} is a subset of Employees. When you call select ["name", "salary"], you’re saying: I have data shaped like the full schema, but I only want the part that fits this smaller schema. The data doesn’t change. You’re just dropping the columns that aren’t in the subset.

What makes this general is the idea of a mapping between schemas. For select, the mapping says “the column name in the small schema corresponds to name in the big schema, and salary corresponds to salary.” That’s an inclusion. For rename "salary" "pay", the mapping says “pay in the new schema corresponds to salary in the old one.” That’s a relabeling. In both cases, there’s a mapping that tells you how to translate between schemas, and the data follows.

Fong and Spivak call this Delta (Δ). Given a mapping between schemas, Delta uses it to reshape the data. It never invents new data or combines rows. It only restructures what’s already there.

You can collapse data along the mapping by merging. Multiple employees share a department. If you want data shaped like the Departments schema (one row per department) you have to decide what to do with all the employees that point at the same department. You could collect them into a list, sum their salaries, or count them. That’s groupBy ["department"] followed by an aggregation.

Fong and Spivak call this Sigma (Σ). Given a mapping where many source rows point at the same target, Sigma collects everything at each target. The collect function keeps all values as a list, which is raw Sigma. Using sum or mean composes the collection with a fold.

You can combine data from both schemas by pairing. Given data on both Employees and Departments, you can find rows that agree on department and stitch them into a wider row. That’s innerJoin ["department"] employees departments. Each result row contains columns from both tables, matched on their shared key.

Fong and Spivak call this Pi (Π). Their mnemonic is “pair and query data,” with a footnote: “more commonly called ‘join’ by database programmers.” Pi finds all tuples that satisfy the constraints imposed by the shared key.

So the three patterns from the previous section (restructuring, merging, pairing) have names: Δ, Σ, Π. But the names are the least interesting part. What’s interesting is why these three arise naturally from the structure of schema mappings.

The answer is about how schemas relate to each other. Fong and Spivak model schemas as categories. A category here is a collection of things (tables, columns) with relationships between them (foreign keys) and a rule for following chains of relationships.

An instance of a schema is what you get when you assign actual data to it. Each table gets a set of rows. Each foreign key gets a function mapping a row to the row it references.

Take the Employees/Departments schema. An instance assigns rows {Alice, Bob, Carol} to Employees, rows {Engineering, Sales} to Departments, and a function that sends Alice → Engineering, Bob → Sales, Carol → Engineering. The only rule is consistency. If Employees references Departments which references Location, looking up Alice’s location directly has to give the same answer as looking up her department and then that department’s location. In category theory, an instance that satisfies this rule is called a functor.

When you have a mapping between two schemas (also a functor), Fong and Spivak prove that it induces these three data migration operations. They’re connected by a structure called an adjoint triple:

Σ ⊣ Δ ⊣ Π

Sigma is the most generous: take everything and merge. Pi is the most conservative: only keep tuples that match on all shared attributes. Delta goes the other direction, restricting data without inventing or combining anything. The adjunction means these three compose cleanly: the output of any Δ step is a valid input for any Σ or Π step, and vice versa. That’s why you can chain select into join into groupBy and the schemas just work out.

What the adjoint triple doesn’t cover

The adjoint triple accounts for five of Petersohn’s seven relational operators. But DIFFERENCE and DROP DUPLICATES, the two I set aside earlier, don’t arise from schema morphisms at all. They operate on instances of the same schema, not on migrations between schemas.

Think about what DIFFERENCE actually does. Suppose you have two dataframes with the same schema:

all_employees:           terminated:
name   department        name   department
Alice  Engineering       Carol  Engineering
Bob    Sales
Carol  Engineering

DIFFERENCE returns the rows in all_employees that don’t appear in terminated: Alice and Bob. The schema doesn’t change. Only the set of rows changes. And DROP DUPLICATES works on a single dataframe:

with_duplicates:         after distinct:
name   department        name   department
Alice  Engineering       Alice  Engineering
Bob    Sales             Bob    Sales
Alice  Engineering

It collapses identical rows into one. Again, the schema is untouched.

The adjoint triple has nothing to say about them because Δ, Σ, and Π are about moving data between schemas. DIFFERENCE and DROP DUPLICATES are about reasoning about subsets of rows within a single schema.

To handle operations on subsets, you need your category to support a notion of “subset” with operations like complement and intersection. Not every category does. But the category of instances on a schema does, because assigning sets of rows to each table and functions to each foreign key gives you all the set-theoretic structure you need. In category theory, a category with this kind of structure is called a topos. The important thing about a topos for our purposes is that it comes equipped with tools for working with subsets:

• DIFFERENCE is a complement operation on subsets. Given two subsets of rows A and B (two dataframes with the same schema), DIFFERENCE computes the rows in A that are not in B. A topos guarantees that this complement is well-defined and behaves the way you’d expect from set theory.

• DROP DUPLICATES is an image operation. Think of each row’s content as a function from the row’s identity to its values. Multiple rows can map to the same values. DROP DUPLICATES collapses them, keeping one representative per distinct value. A topos guarantees that every such function factors cleanly through its image, which is what makes DROP DUPLICATES well-defined.

So the categorical picture of relational operations has two layers: the migration functors (Δ, Σ, Π) for moving data across schemas, and the topos structure for reasoning about subsets of rows within a schema.

Here’s the full picture, mapping Petersohn’s operators to the categorical framework:

The first five operators decompose into the adjoint triple Δ, Σ, Π. The next two use the topos structure of the category of instances. Together, these seven account for the relational core. The next three preserve the schema. They’re morphisms between instances on the same schema, not migrations between schemas. The last four are outside the relational model entirely.

This is the compression: 200 pandas operators → 15 algebraic operators → 3 migration functors and 2 topos-theoretic operations covering the relational core. The schema-preserving and dataframe-specific operators are important, but they’re not where the complexity lives. The migration functors handle schema-changing operations, and the topos structure handles set-theoretic reasoning within a schema.

Designing an API around these patterns

The categorical decomposition gives you a design principle: each operation should have a clear rule for computing its output schema from its input schema. The migration functors change the schema. The topos-theoretic operations preserve the schema but change which rows are present.

Start with the migration functors. Delta gives you operations that reshape a schema without computing new data. That means:

• select columns → output schema is the subset of input columns you named
• exclude columns → output schema is the input minus the columns you named
• rename old new → output schema is the input with one column relabeled

These operations share a property: given the input schema and the operation’s arguments, you can compute the output schema without looking at any data. That makes them cheap, predictable, and safe to reorder in an optimizer.

Then Sigma. You need operations that collapse rows by key. That means:

• groupBy keys followed by aggregate [sum ..., mean ..., count ...] → output schema is the key columns plus one new column per aggregation
• groupBy keys followed by collect → output schema is the key columns plus list-valued columns

The key insight from the categorical picture is that collect and aggregate are variations of the same pattern. Sigma collects everything at each key, and aggregation functions like sum or mean are an optional reduction on top.

Then Pi. You need operations that combine two schemas by shared key. That means:

• innerJoin keys left right → output schema is key columns (once) plus non-key columns from both sides
• leftJoin keys left right → same, but right-side non-key columns become nullable
• fullOuterJoin keys left right → same, but non-key columns from both sides become nullable

Each join variant is Pi with a different policy for missing matches. The schema rule is the same; only the nullability wrapping changes.

Then the topos layer. DIFFERENCE and DROP DUPLICATES don’t need schema rules because they preserve the schema. Their type signatures reflect this:

distinct :: TypedDataFrame cols -> TypedDataFrame cols
-- difference :: TypedDataFrame cols -> TypedDataFrame cols -> TypedDataFrame cols

The input and output types are identical. What changes is which rows are present: distinct collapses duplicates, and difference removes rows that appear in the second argument. Both take a dataframe and return a dataframe with the same columns but fewer rows, which makes their semantics straightforward to implement and optimize.

Finally, the schema-preserving operations (filter, sort, take, sample) sit outside the migration and topos patterns. Their output schema is always identical to their input schema. They’re important, but they don’t interact with schema composition, which is why they can be designed independently.

Once you have these pieces, a pipeline is a chain of migration steps (Δ, Σ, Π) and row-level steps (DIFFERENCE, DROP DUPLICATES, filter), and each step’s output schema is a valid input for the next. In the dataframe library, Haskell’s type system enforces this. Schemas are encoded at the type level, and the compiler checks every transition:

type Employees =
    '[ Column "name" Text
     , Column "department" Text
     , Column "salary" Double
     ]

type Departments = '[ Column "department" Text, Column "budget" Double ]

result =
    employees
        & T.distinct                                -- topos: drop duplicate rows
        & T.innerJoin @'["department"] departments  -- Π: schema grows
        & T.derive @"cost_ratio"                    -- grows by one
            (T.col @"salary" / T.col @"budget")
        & T.select @'["department", "cost_ratio"]   -- Δ: schema shrinks
        & T.groupBy @'["department"]
        & T.aggregate                               -- Σ: collapse
            ( T.agg @"avg_ratio" (T.mean (T.col @"cost_ratio"))
            $ T.aggNil
            )

Reference "salary" after select drops it? Compile error. Derive a column that already exists? Compile error. Join on a key missing from one table? Compile error. If the pipeline compiles, every schema transition is valid. This isn’t Haskell-specific. Any language with sufficiently expressive types could enforce the same rules. The categorical decomposition tells you what the rules are; the type system enforces them.

The patterns also help with optimization. If you know exactly what each operation does to the schema, you can reason about when it’s safe to reorder steps. The library has a lazy evaluation mode where a pipeline is built up as a logical plan before being executed:

optimize :: Int -> LogicalPlan -> PhysicalPlan
optimize batchSz =
    toPhysical batchSz
        . eliminateDeadColumns
        . pushPredicates
        . fuseFilters

Consecutive filters get fused into one. Filters get pushed past column operations toward the data source. Derived columns that are never referenced downstream get dropped before execution. These rewrites are safe because the operations obey algebraic laws that follow from the categorical structure: restructuring and filtering commute when the filter doesn’t touch restructured columns; conjunction of predicates is the same as filtering twice; a Δ step that drops a column can’t affect a filter that doesn’t reference it.

Where this is going

What I’m working toward is a canonical definition of the dataframe. Petersohn et al. made the best attempt I’ve seen with their data model and algebra. Category theory adds structure on top: three migration functors for schema-changing operations, and topos structure for set-theoretic reasoning within a schema. Together, these cover the relational core.

That’s what I wanted when I started the library. A small set of operations grounded in theory, with the compiler verifying every step. This post covers the relational operators, the ones dataframes share with SQL. The dataframe-specific operators (TRANSPOSE, TOLABELS, FROMLABELS) and the symmetry between rows and columns deserve their own treatment. But for the relational core, the two-layer picture works.

If any of this sounds interesting: Fong and Spivak’s Seven Sketches in Compositionality is written for non-mathematicians and builds from first principles. Chapter 3 covers databases and the Δ/Σ/Π migration functors. Petersohn et al.’s Towards Scalable Dataframe Systems covers the algebra, the data model, and what people actually do in those 1 million notebooks. Both are worth your time.

The dataframe library is on GitHub. The typed API lives in DataFrame.Typed. I’d love to hear what you build with it.

If you’ve done this enough times, you start to recognize the repeating shapes: “per unit,” “per time,” “per person,” “net,” “remaining,” “rate.” These show up everywhere because they match how the world is measured.

The part I’ve been curious about is: how much of this can we automate without losing that “nameable quantity” feeling?

This post is about an experiment that worked surprisingly well for me. I treat feature engineering as a tiny program synthesis problem: enumerate candidate derived features as small arithmetic expressions over columns. Then I train a decision tree using rules built from those derived features. The twist is how I keep the feature generator from spiraling into unintelligible math: I use an LLM to prune out feature combinations that don’t make sense.

The LLMs isn’t the model making the modelling decisions. The LLM’s job is much smaller: look at a candidate arithmetic expression and answer one question—does this resemble a meaningful real-world quantity a human might compute?

That one constraint was enough to turn a messy tree into something that looks a lot more like human feature engineering.

You can clone the repo and try it out yourself.

Features and meaning

In practice, the strongest features are the ones you can explain simply. To yourself and to others.

If I write down clicks / impressions, I can call it “conversion rate.” If I write down price / sqft, I can call it “price per square foot.” If I write down revenue - costs, I can call it “profit.”

Even when those aren’t perfect, they’re at least coherent. They’re quantities that exist in the world.

Now compare that to something like income / zip_code. It might correlate with house price (New York incomes are higher and would be scaled down less because they have smaller valued zip codes compared to Texas, for example). It might even boost accuracy. But it’s not a quantity anyone would compute. It’s a red flag that the search procedure is exploiting quirks in the dataset instead of discovering “features.”

This distinction—between “works statistically” and “is a coherent quantity” is exactly what I wanted to encode into the feature generator.

The demo dataset

I used Kaggle’s Titanic dataset because it’s familiar and because there’s a well-known “starter pack” of feature engineering that people do. In particular, extracting a title from the passenger name (“Mr.”, “Mrs.”, etc.), deriving a cabin prefix (the leading letter), and doing something with ticket prefixes are all common moves.

I wanted this demo to feel honest: do the obvious parsing that many Titanic solutions do, then stop. Let the system take it from there.

The full pipeline I ran is in this file.

A couple things are doing quiet work here.

First: I’m renaming columns so the output reads like English. That matters more than it sounds. If the whole point is interpretability, the printed artifact shouldn’t feel like you’re decoding abbreviations.

Second: notice maxExprDepth = 2. That’s the “I want features I can still name” constraint. Depth 3 can be fun, but depth 2 is usually where you get ratios, differences, and simple interactions without the expression turning into algebra homework.

Third: I set complexityPenalty = 0 in this run on purpose. I wanted to isolate the effect of the LLM filter. In other words: “what happens if the only regularization pressure is semantic?”

The feature generator

The mechanics are simple. I generate candidate numeric expressions from the columns. Things like a + b, a / b, a - b, a * b, plus the raw columns themselves. Then I turn each numeric expression into candidate rules by picking thresholds at a small grid of percentiles.

At that point, tree learning is as boring as you want it to be: pick the rule that best improves impurity, split, recurse, stop at depth/min-sample constraints. For classification trees, using something like Gini impurity is standard practice.

So far, this is just automated feature generation feeding a greedy tree. The entire problem is that the expression generator will happily propose nonsense.

The problem

When you enumerate expressions naively, you’ll get candidates like Fare + Age. Strange combinations of class codes and counts that are hard to defend. None of this is illegal. Some of it is even predictive. But it’s not the kind of feature engineering that’s robust or you’d want to explain to someone else. It’s not the kind of artifact you’d want to ship into a risk-sensitive workflow.

This is a classic program synthesis story in miniature: enumerative search grows quickly, and most of the space isn’t interesting. Priors and pruning are the difference between “works” and “does anything you want.”

In my case, the “interesting” region of the search space is: expressions that resemble real derived quantities. Which is a semantic judgment that is awkward to encode by hand. So I outsourced that judgment to a very constrained LLM call.

The trick: use an LLM as a semantic regularizer

For each candidate numeric expression, I ask an LLM to output a single digit score from 0 to 10. The prompt includes a short rubric: addition and subtraction should combine like-with-like; multiplication and division can produce useful derived quantities; categorical codes should be treated with suspicion; conversion functions don’t matter.

Then I do the most unglamorous thing possible: I drop expressions that score below a threshold.

That’s it. I’m not asking the LLM to invent features. I’m not asking it to choose thresholds. I’m not asking it to “solve Titanic.” I’m using it like a regularizer: a pressure that biases what the learner is allowed to consider.

This “LLM as guidance for search” framing is showing up more and more in synthesis-style work: use the model to bias exploration, propose edits, or filter candidates, but keep a separate objective that actually evaluates what works.

I like the mental model: the LLM is the bouncer.

The before/after that made me confident this was worth writing down I ran the exact same pipeline twice.

In the first run, I didn’t apply the semantic filter. I just let the candidate generator dump its expression space into the tree learner and let impurity gain decide.

This is the “without LLM” tree I got (translated from the DSL):

  ifThenElse
    (eq (col @Maybe Text "title") (lit (Just "Mr.")))
    
    THEN:
    ├─ ifThenElse
    │    (or
    │      (geq (toDouble (col @Int "passenger_class")) (lit (2.0)))
    │      (eq (col @Maybe Text "cabin_prefix") (lit (Nothing)))
    │    )
    │    
    │    THEN:
    │    ├─ lit (0)
    │    
    │    ELSE:
    │    └─ ifThenElse
    │         (lt
    │           (ifThenElse
    │             (eq (toDouble (col @Int "number_of_siblings_and_spouses")) (lit (0.0)))
    │             THEN: divide
    │                     (toDouble (col @Int "passenger_class"))
    │                     (toDouble (col @Int "number_of_siblings_and_spouses"))
    │             ELSE: lit (0.0)
    │           )
    │           (lit (1.0))
    │         )
    │         
    │         THEN: lit (0)
    │         ELSE: lit (1)
    
    ELSE:
    └─ ifThenElse
         (lt (toDouble (col @Int "passenger_class")) (lit (3.0)))
         
         THEN:
         ├─ lit (1)
         
         ELSE:
         └─ ifThenElse
              (gt
                (mult
                  (toDouble (col @Int "number_of_parents_and_children"))
                  (toDouble (col @Int "number_of_siblings_and_spouses"))
                )
                (lit (1.0))
              )
              
              THEN:
              ├─ lit (0)
              
              ELSE:
              └─ ifThenElse
                   (and
                     (eq (col @Maybe Text "Embarked") (lit (Just "S")))
                     (geq (col @Maybe Double "Age") (lit (Just 14.650000000000002)))
                   )
                   
                   THEN:
                   ├─ ifThenElse
                   │    (gt
                   │      (ifThenElse
                   │        (eq (toDouble (col @Int "number_of_siblings_and_spouses")) (lit (0.0)))
                   │        THEN: divide
                   │                (toDouble (col @Int "passenger_class"))
                   │                (toDouble (col @Int "number_of_siblings_and_spouses"))
                   │        ELSE: lit (0.0)
                   │      )
                   │      (lit (1.0))
                   │    )
                   │    THEN: lit (0)
                   │    ELSE: lit (1)
                   
                   ELSE:
                   └─ lit (1)

The DSL is in prefix/Polish notation. So (eq (col @Maybe Text "title") (lit (Just "Mr."))) means title == "Mr.". The Maybe part is just how Haskell deals with nulls.

Even if you don’t read every branch, you can feel the flavor. There are lots of conditions that look like dataset artifacts: dividing passenger class by number of siblings, multiplying number of parents with number of siblings. This is exactly what happens when the search is allowed to treat “anything correlated” as a legitimate feature.

The training accuracy and validation of this decision tree were 0.82.

What happens when we use the LLM?

Getting the prompt right

I started off with this prompt:

Please return nothing else but a single number between 0 and 10 inclusive denoting how interpretable the following expression is and if the operation makes sense given your inferred dimensions of the variables. 0 would be something completely meaningless e.g rgb_house_colour * lot_size and 10 would be something that is a commonly reasoned about quantity e.g price / lot_size. The expression is: fare + age

Which was admittedly confusingly written and underspecified. After trying it a number of times in the ollama CLI and getting counterintuitive results (number_of_people + rejection_rate scored high but number_of_siblings + number_of_parents scored low), I iterated on the prompt to more thoroughly specify the task. Since ollama doesn’t “think” you have to be careful how you engineer the prompt so it gets a reasonable solution in one shot.

After many iterations I settled with:

You must output ONLY a single digit 0-10, nothing else. No explanation, no text, just the number.

Evaluate if this expression produces a meaningful quantity by checking:
    - Do the units/types match for the operation?
    - Is either operand a categorical code rather than a true quantity?
    - Would the result be useful to actually calculate?

Scoring:
    - 0-3: Result is meaningless (e.g., "fare + age" = dollars+years, "height + weight" = meters+kg)
    - 4-5: Unclear or context-dependent meaning
    - 6-7: Makes sense in specific domains
    - 8-9: Clear, commonly useful quantity
    - 10: Fundamental/universal quantity
Guidelines:
    - Addition/subtraction: operands must represent the same kind of thing
    - Multiplication/division: can create meaningful derived quantities
    - Consider: would this result be useful to calculate in practice?
    - `toDouble` is just a function that converts any number to a decimal and is semantically unimportant.

Examples:
    toDouble(fare) + toDouble(age) = 2 (adding money to years)
    toDouble(price) / toDouble(area) = 9 (price per sq ft)
    toDouble(distance) / toDouble(time) = 10 (speed)
    toDouble(num_people) * toDouble(rejection_rate) = 9 (expected rejections)
    toDouble(revenue) - toDouble(costs) = 10 (profit)
    toDouble(height) + toDouble(weight) = 2 (adding length to mass)
Output format: Just the digit, e.g., 2
Think very carefully about each but only give me the final answer.
Expression:

Here’s the second tree where candidate expressions are pruned by the LLM if their “semantic” score is less than 5:

model =
  ifThenElse
    (and
      (lt (toDouble (col @Int "ticket_class")) (lit (3.0)))
      (eq (col @Text "Sex") (lit ("female")))
    )
    
    THEN:
    ├─ lit (1)
    
    ELSE:
    └─ ifThenElse
         (and
           (eq (col @Maybe Text "cabin_prefix") (lit (Nothing)))
           (eq (col @Maybe Text "title") (lit (Just "Mr.")))
         )
         
         THEN:
         ├─ lit (0)
         
         ELSE:
         └─ ifThenElse
              (geq
                (add
                  (toDouble (col @Int "number_of_parents_and_children_aboard"))
                  (toDouble (col @Int "number_of_siblings_and_spouses_aboard"))
                )
                (lit (3.0))
              )
              
              THEN:
              ├─ lit (0)
              
              ELSE:
              └─ ifThenElse
                   (or
                     (gt (toDouble (col @Int "ticket_class")) (lit (1.0)))
                     (geq (col @Maybe Double "Age") (lit (Just 14)))
                   )
                   
                   THEN:
                   ├─ lit (0)
                   
                   ELSE:
                   └─ lit (1)

The initial condition is already immediately readable: women in first or second class survived. Even deeper into the tree it remains readable - large faimiles didn’t survive (encoded in the conditon: if number of parents and siblings is greater than 3 then 0).

The training and validation accuracy of this decision tree were both 0.83. Better than the first model!

More importantly the tree was simpler. The change wasn’t subtle. The tree got smaller and more coherent. It rediscovered features that are basically “human Titanic features” like family size, and simple interactions with sex. In other words: it started behaving like a feature engineering notebook, not like a correlation-miner.

This is the moment the “LLM as regularizer” framing really clicked for me. The learner didn’t change. The objective didn’t change. The only thing that changed was the hypothesis space: which derived quantities were allowed to exist as candidates.

And that was enough to nudge the output from “math that correlates” to “features you can name.”

If you’re reading this and thinking “Titanic is toy-ish,” I agree. I used it because everyone already knows what reasonable features look like, which makes the interpretability comparison crisp. The real test is whether this behavior carries over to messier tabular problems where feature engineering actually hurts.

Why I think this is a nice way to use LLMs

There are two reasons I like this pattern.

First, it doesn’t throw the entire problem at the LLM hoping it gets it right. The LLM is asked to do semantic triage. That’s a smaller, more stable job.

Second, it treats interpretability as a first-class constraint. If your model is going to be read by humans, you don’t want to optimize only for accuracy and then bolt on interpretability later. You want the candidate language itself to be biased toward human-shaped quantities.

What I have right now is still a prototype, but the shape feels right: a small synthesis loop producing candidate features, a classic learner picking among them, and an LLM providing a soft semantic prior when you don’t have explicit unit metadata.

Limitations (the kind you notice immediately)

The first limitation is determinism. If you want this to be reproducible, you’ll want caching: expression string → score. You’ll also want deterministic decoding, or at least stable settings.

The second limitation is schema context. The rubric can only be as good as the column names. If the dataset has columns like X1, X2, A, B, you shouldn’t expect semantic miracles. The filter can become much better if you pass minimal schema descriptions (“this is dollars,” “this is years,” “this is a categorical ID”).

The third limitation is that “meaningful quantity” is subjective. There are domains where “Age <= Fare” might encode something real (it’s still weird, but you get the idea). This is why I think of the score as a regularizer rather than a hard rule: it’s a knob you tune based on how strict you want the feature language to be.

What I want to try next

The next steps feel straightforward.

One is to distill the LLM into something cheaper. Once you have a cache of labeled expressions, you can train a small local classifier that approximates the semantic score and only calls the LLM for uncertain cases.

Another is to combine semantic regularization with structural regularization. In the demo run above I set the explicit complexity penalty to zero to isolate the LLM effect. In practice, you can (and probably should) use both: penalize expression size while also filtering out expressions that fail a semantic sanity check.

And the most exciting direction, for me, is to apply this to real tabular workflows where the “right derived quantities” aren’t obvious, but still exist: fraud, risk, forecasting, ops metrics, churn, all the places where humans already reason in rates, margins, and per-unit measures.

If nothing else, this experiment made one thing feel concrete: there’s a useful middle ground between “manual feature engineering forever” and “opaque feature learning.” You can automate the search, and still keep the artifact readable, if you treat semantics as something the search needs to respect.

This post kicks off the second part of a hands-on series about program synthesis. We’ll apply the previously explored technique (an enumerative bottom-up search) to a slightly more realistic problem: automatically generating features for the Iris dataset.

A recap

In the last post, we created a Flash-fill-style system that generates programs that transform strings given input-output examples as specifications. We defined a domain-specific language and searched the space of all possible programs to find one that produced all the inputs for our desired outputs.

This search technique (enumerative bottom-up search) is effectively a breadth-first search. The high-level algorithm is as follows:

search(inputs, outputs) {
 programs = all_simple_programs()
    while (true) {
        for (program : programs) {
            if (satisfiesExamples(inputs, outputs, program)) {
                return program
 }
 }
 programs = deduplicate(inputs, expand(programs))
 }
    // no program was found
    return null
}

The limits of this approach

Combinatorial explosion

We tried to deal with the combinatorial explosion by deduplicating and keeping our DSL small. While effective for our toy problem, most real-world problems aggressively resist attempts to be so easily constrained. Interesting problems typically require expressive DSLs (binary and ternary functions, recursive functions, conditional logic, etc.), which make searching much more difficult. Deduplicating expressions also gets difficult as the DSL becomes more expressive.

We’ll have to figure out a way to tame the algorithmic complexity of the problem without hamstringing our search.

What is a correct program?

Checking if a program satisfies our examples works well when we can trust that the examples:

• are complete and representative of the problem
• correctly specify the underlying program we are trying to discover.

When dealing with data in the wild, we can’t always guarantee that our inputs or outputs are correct. Our data is typically a noisy sub-sample of a larger distribution. Finding the exact program that satisfies a given set of examples might create an inflexible, complicated program that tries to get every single example right. In mathematical modelling, we call this overfitting: when a program fits the training examples but doesn’t generalise well.

We’re going to have to come up with approximate notions of “correctness” that can let us know when we’ve found a program that’s good enough for the problem.

Our problem

In a previous post, we trained a simple neural network to classify plants into three iris species. With some very light, hand-wavy feature engineering, we did well.

Our test set performance wasn’t as good as we’d hoped.

Our challenge for this post will be to engineer better features using symbolic regression.

Wait…symbolic regression? What happened to program synthesis? Well, symbolic regression is a kind of program synthesis in which the programs are mathematical expressions, and instead of finding the exact program, we find one that fits a dataset. The underlying algorithm is practically the same:

search(inputs, outputs, depth) {
 formulas = all_simple_formulas()
    while (depth > 0) {
 forumlas = deduplicate(inputs, expand(formulas))
 }
    return pickBestFormula(inputs, outputs, formulas)
}

There are a couple of significant differences between this and our earlier synthesis algorithm.

• Our target is uncertain, so we have to constrain our search to only work until a certain depth.
• Because we are looking for a function that best approximates the relationship between the input and output, we have to include some logic to pick the best solution (for some definition of best).

This process should yield features we can directly plug into our neural network.

A prototype of this approach exists in the dataframe library, so this post will explain that code. I have simplified the code for this post.

Why feature engineering?

Feature engineering is a more straightforward target than finding an exact program (we’d be building a fully fledged classifier). The result of feature engineering must have a strong relationship with the target variable. We’ll see some tools for finding exact expressions later in this post.

Our search space

The dataframe library defines an expression DSL that looks roughly like this:

data Expr a where
 Col :: String -> Expr a  -- A reference to a column in the dataframe
 Lit :: a -> Expr a       -- A literal value
 UnaryOp :: (b -> a) -> Expr b -> Expr a
 BinaryOp :: (c -> b -> a) -> Expr c -> Expr b -> Expr a

This DSL is effectively the same expression language we defined in rewriting dataframes for MicroHs.

We can define some mathematical operations using this DSL.

add :: (Num a) => Expr a -> Expr a -> Expr a
add = BinaryOp (+)

sub :: (Num a) => Expr a -> Expr a -> Expr a
sub = BinaryOp (-)

cos :: Num a => Expr a -> Expr a
cos = UnaryOp Prelude.cos

...

These will be our simple programs/formulas (from hereon, I’ll use the two words interchangeably), which we use to seed and grow the search space.

What we call the expand function in the pseudo code corresponds to the generatePrograms function in the actual implementation. We define it as such:

generatePrograms ::
 [Expr Double] ->
 [Expr Double] ->
 [Expr Double] ->
 [Expr Double]
generatePrograms vars constants [] = vars ++ constants
generatePrograms vars constants ps =
    let
 existingPrograms = ps ++ vars ++ constants
     in
 existingPrograms
 ++ [ transform p
 | p <- ps ++ vars
 , transform <-
 [ sqrt
 , abs
 , log . (+ Lit 1)
 , exp
 , sin
 , cos
 , relu
 , signum
 ]
 ]
 ++ [ pow i p
 | p <- existingPrograms
 , i <- [2 .. 6]
 ]
 ++ [ p + q
 | (i, p) <- zip [0 ..] existingPrograms
 , (j, q) <- zip [0 ..] existingPrograms
 , i >= j
 ]
 ++ [ p - q
 | (i, p) <- zip [0 ..] existingPrograms
 , (j, q) <- zip [0 ..] existingPrograms
 , i /= j
 ]
 ++ [ p * q
 | (i, p) <- zip [0 ..] existingPrograms
 , (j, q) <- zip [0 ..] existingPrograms
 , i >= j
 ]
 ++ [ p / q
 | p <- existingPrograms
 , q <- existingPrograms
 , p /= q
 ]

At each iteration (except the first), we take the list of existing formulas and compose them with simple programs. For example, when we invoke generatePrograms [] [] [col "x"] we get sqrt(x), abs(x),..., pow 0 x, pow 2 x,...,x + x, x * x.

We’ve expanded our initial “formula” into a set of prospective formulas. We already notice some duplicates here. pow 2 x == x * x. To save ourselves the trouble of an expansion branch, we should keep only one of these.

Deduplication

Mathematics has many tools for simplifying complex structures into canonical forms. For example, 1/3, 2/6, and 4/12 can effectively be treated as 1/3. Until now, we have relied on observational equality to deduplicate expressions. Such a deduplication would quickly catch a fair number of these deduplications, but it’s still a relatively expensive notion of equality - especially as the number of formulas grows. We need a way to leverage our regular tools for mathematical equivalence so we don’t always have to evaluate our prospective formulas on the dataframe.

Luckily, this is a solved problem. Symbolic regression implementations typically use equality graphs (e-graphs) to rewrite expressions in their simplest form. E-graphs are a valuable tool in compilers as well, since they can help optimise large expressions into simpler equivalents.

In this example, however, we won’t use equality graphs. For pedagogical reasons, we’ll implement our own custom rewrite rules.

We can do this in two ways:

• Bake reduction rules into the expressions of the DSL
• Normalise whole expressions into their canonical forms post-hoc.

(+) :: Expr a -> Expr a -> Expr a
(+) (Lit x) (Lit y) = Lit (x + y)
(+) e1 e2
    -- Simple rule that says if we're adding x to itself then rewrite the expression as x * 2
 | e1 == e2 = UnaryOp ("mult " <> (T.pack . show) (Lit @a 2)) (* 2) e1
 | otherwise = BinaryOp "add" (+) e1 e2


normalize :: (Eq a, Ord a) => Expr a -> Expr a
normalize expr = case expr of
 Col name -> Col name
 Lit val -> Lit val
 If cond th el -> If (normalize cond) (normalize th) (normalize el)
 UnaryOp name f e -> UnaryOp name f (normalize e)
 BinaryOp name f e1 e2
 | isCommutative name ->
            let n1 = normalize e1
 n2 = normalize e2
             in if compareExpr n1 n2 == GT
                    then BinaryOp name f n2 n1 -- Swap to canonical order
                    else BinaryOp name f n1 n2
 | otherwise -> BinaryOp name f (normalize e1) (normalize e2)

Of course, these are straightforward optimisations. An e-graph would implement many smart rewrite rules, but we’ll stick with these for now.

The deduplication function can thus be written as:

deduplicate ::
    DataFrame ->
 [Expr Double] ->
 [(Expr Double, TypedColumn Double)]
deduplicate df = go S.empty . nubOrd . L.sortBy (\e1 e2 -> compare (eSize e1) (eSize e2)) . map normalize
  where
 go _ [] = []
 go seen (x : xs)
 | hasInvalid = go seen xs
 | S.member res seen = go seen xs
 | otherwise = (x, res) : go (S.insert res seen) xs
      where
 res = either throw id (interpret @Double df x)
        -- We need this since we allow functions that might return infinity or NaN
 hasInvalid = VU.any (\n -> isNaN n || isInfinite n) (toVector res)

Preventing combinatorial explosion

Deduplication isn’t enough. We need more creative tools to combat the combinatorial explosion as we explore the search space. Most of the ingenuity of enumerative search techniques lies here.

We have a couple of options here:

• Greedy: take the n best expressions each round (otherwise called a beam search)
• Divide-and-conquer: first train a neural network to use as an oracle to guide your search, then run tests against the neural network to attempt to deconstruct the function (this is called the AI Feynman algorithm)
• Evolution: randomly mutate and mix functions to better explore the space of programs (genetic programming)

For simplicity, we’ll go with the greedy approach. Most symbolic regression tools use genetic programming, but that approach would be too powerful for our problem.

The greedy approach entails evaluating the expression on some data (which we’ve already done and returned from the deduplicate function), then picking the best n functions. We’ll use simple Pearson’s correlation to rank our formulas.

pickTopN ::
    DataFrame ->
    TypedColumn Double ->
    BeamConfig ->
 [(Expr Double, TypedColumn a)] ->
 [Expr Double]
pickTopN _ _ _ [] = []
pickTopN df (TColumn col) cfg ps =
    let
 l = toVector @Double col
 ordered =
            Prelude.take
 (beamLength cfg) -- 4. Take only the first n
                L.sortBy
 ((flip compare) `on` snd) -- 3. Sort functions by descending objective.
 (map
 (\(e, res) -> (e, getLossFunction (lossFunction cfg) l (asDoubleVector res)))  -- 2. Compute objective
 ps -- 1. programs
 )
     in map fst ordered

Now we have everything we need to define our beam search:

beamSearch ::
    DataFrame ->
 -- | Parameters of the beam search.
    BeamConfig ->
 -- | example labels
    TypedColumn Double ->
 -- | Constants
 [Expr Double] ->
 -- | Conditions
 [Expr Bool] ->
 -- | Programs
 [Expr Double] ->
    Maybe (Expr Double)
beamSearch df cfg labels constants conds programs
 | searchDepth cfg == 0 = case ps of
        [] -> Nothing
 (x : _) -> Just x
 | otherwise =
 beamSearch
 df
 (cfg{searchDepth = searchDepth cfg - 1})
 labels
 constants
 (generatePrograms vars constants ps)
  where
 vars = map col names
 ps = pickTopN df labels cfg $ deduplicate df programs
 names = D.columnNames df

Retraining iris

We create a training split to use for our feature engineering. Doing feature engineering on a training set avoids leaking test-time information.

import qualified DataFrame.Hasktorch as DHT
import qualified System.Random as SysRand
import DataFrame.Functions ((.==))

let (training, _) = D.randomSplit (SysRand.mkStdGen 12345) 0.7 df

We then create a separate feature for each class by setting its instances to 1 and everything else to 0.

labelWith :: Text -> DataFrame -> DataFrame
labelWith species = training |> D.derive "variety" (F.ifThenElse (variety .== species) (F.lit @Double 1) 0)

versicolor = F.synthesizeFeatureExpr "variety" F.defaultBeamConfig (labelWith "Versicolor")
virginica = F.synthesizeFeatureExpr "variety" F.defaultBeamConfig (labelWith "Virginica")
setosa = F.synthesizeFeatureExpr "variety" F.defaultBeamConfig (labelWith "Setosa")

Running our synthesizeFeatureExpr function gives us the following features:

• virginica: (col @Double "petal.length") * (col @Double "petal.width")
• versicolor: cos (col @Double "petal.length")
• setosa: 1 / (col @Double "petal.length")

Let’s look through each of these features one by one.

Virginica

This feature calculates the approximate petal area. Virginica irises have the largest petals, so this product yields the highest values for this species. For example, if virginica has petals ~6cm long and ~2cm wide, this gives ~12, while smaller-petaled species produce smaller values.

Versicolor

This feature applies a cosine transformation to petal length. This is the most clever/unusual one. Since:

• Setosa has short petals (~1.0-1.9 cm)
• Versicolor has medium petals (~3.0-5.1 cm)
• Virginica has long petals (~4.5-6.9 cm)

The cosine function’s oscillating nature means that versicolor’s petal lengths (around 3-5 radians) fall into a specific range of the cosine curve that differs from the other species. This creates separation even though versicolor is “in the middle” on raw measurements.

This suggests that the search was trying to model some non-linearity which could be approximated by cos.

Setosa

This feature creates an inverse relationship with petal length. Since setosa has the shortest petals (~1-2 cm), dividing 1 by small numbers yields large values (e.g., 1/1.5 ≈ 0.66). Longer-petaled species get much smaller values (e.g., 1/6.0 ≈ 0.16).

Initial results

What happens when we feed these features into the neural network? Well, we get better at the training set:

But the test set is a mixed bag.

We do better at classifying setosas, but we still have trouble differentiating between versicolor and virginica. The cos must have overfit to the non-linear pattern in the training set but failed to generalise. How can we better model the separation between versicolor and virginica?

Let’s throw some conditionals into our DSL and see how far that gets us.

Conditionals are your frenemy

Conditionals model many valuable programs. Armed with conditionals, we would be able to express almost any function. That’s a good thing, but it also means we can even express the meaningless functions that arbitrarily overfit our data. Conditionals also mean our search space blows up. Take a simple if-statement that always checks <= in the condition. if p <= q then r else s. Computing this for all existing programs does O(n⁴) at each tree level. This brings our search to a grinding halt at just the second level, even with a small dataset and feature space.

How can we make conditionals tractable? Well, we can run a beam search over just the conditionals, combining them to form conditions that are strongly related to the target variable. Because we want to capture interesting non-linearities, conditionals are evaluated, converted to 0 or 1, then ranked by mutual information rather than Pearson correlation.

generateConditions ::
    TypedColumn Double -> [Expr Bool] -> [Expr Double] -> DataFrame -> [Expr Bool]
generateConditions labels conds ps df =
    let
 newConds =
 [ p .<= q
 | p <- ps
 , q <- ps
 , p /= q
 ]
 ++ [ DataFrame.Functions.not p
 | p <- conds
 ]
 expandedConds =
 conds
 ++ newConds
 ++ [p .&& q | p <- newConds, q <- conds, p /= q]
 ++ [p .|| q | p <- newConds, q <- conds, p /= q]
     in
 pickTopN 10 df labels (deduplicate df expandedConds)

Now, our program supports two expansions—one for conditionals and another for programs.

generatePrograms _ _ vars' constants [] = vars' ++ constants
generatePrograms includeConds conds vars constants ps =
    let
 existingPrograms = ps ++ vars ++ constants
     in
 ...
 ++ ( if includeConds
                    then
 [ DataFrame.Functions.min p q
 | (i, p) <- zip [0 ..] existingPrograms
 , (j, q) <- zip [0 ..] existingPrograms
 , Prelude.not (isLiteral p && isLiteral q)
 , p /= q
 , i > j
 ]
 ++ [ DataFrame.Functions.max p q
 | (i, p) <- zip [0 ..] existingPrograms
 , (j, q) <- zip [0 ..] existingPrograms
 , Prelude.not (isLiteral p && isLiteral q)
 , p /= q
 , i > j
 ]
 ++ [ ifThenElse cond r s
 | cond <- conds
 , r <- existingPrograms
 , s <- existingPrograms
 , r /= s
 ]
                    else []
 )

Generated features

Running our synthesizeFeatureExpr function with conditionals enabled gives us the following features:

• virginica: ifThenElse (or (leq (col @Double "petal.length") (col @Double "sepal.width")) (leq (lit (1.8)) (col @Double "petal.width"))) (lit (7.7)) (col @Double "petal.length")
• versicolor: ifThenElse (or (leq (lit (5.1)) (col @Double "petal.length")) (leq (lit (1.8)) (col @Double "petal.width"))) (lit 1) (lit 0)
• setosa: ifThenElse (leq (col @Double "sepal.width") (col @Double "petal.length")) (lit 0) (lit 1)

Virginica

This feature creates a derived feature that:

• Returns a constant value of 7.7 when either the petal is unusually short compared to sepal width, OR the petal width is ≥ 1.8 (a wide petal)
• Otherwise returns the actual petal length
• This likely captures virginica characteristics (larger petals)

Versicolor

This feature is a binary indicator that flags:

• Flowers with petal length ≥ 5.1 OR petal width ≥ 1.8
• Returns 1 (true) or 0 (false)
• This targets medium-sized flowers typical of versicolor

Setosa

This feature is a binary indicator that returns 1 when:

• Sepal width <= petal length
• Setosas typically have short petals relative to sepal width, so this would return 1 for setosas and 0 for others.

Final results

With these 3 features, we get almost everything right on the training set.

Our performance on the test set has improved and is much more balanced.

This is a clear improvement over both the previous versions.

We did it! We used interpretable, automatically generated features to improve our neural network!

The symbolic regression landscape

Many battle-tested tools employ the techniques described here. PySr is the most mature in the Julia/Python ecosystem. Haskell has SRTree (which we plan to use eventually in the dataframe library).

Wanna learn more about bottom-up synthesis? MIT offers an introduction to program synthesis, with a problem set that roughly covers what we’ve done in these last two posts. In fact, I encourage you to audit the class to keep abreast with what’s happening in the neuro-symbolic world.

What’s next?

Conditionals were a powerful tool, but posed a big challenge for us. How do we use them without blowing up the search space? What if we could tell the search the shape of the final program (that it’s a conditional with 2 or 3 terms, and it returns 1 or 0)? We would only search for the hole we need to fill, not worry about others searching everything else.

Top-down enumerative search will be the subject of our next post in the series.

Wanna join in this work?

Come over to the dataHaskell Discord where we discuss the work that is at the intersection of functional programming and data science?

At a high level the intent of this post is to:

• get people excited about how far the ecosystem has come,
• checkpoint what the ecosystem currently looks like,
• identify areas of improvement in the workflow.

This post is a narrative companion to a hosted notebook, but with a little more context.

This work was done as part of the dataHaskell: an organization devoted to enabling reliable and reproducible machine learning using the Haskell programming language.

Special thanks to Jireh Tan for writing the code.

The iris classification problem

The Iris dataset is machine learning’s “Hello, World!” created by statistician Ronald Fisher in 1936. The dataset contains measurements of 150 iris flowers from three species:

• Setosa
• Versicolor
• Virginica

For each flower, we have four measurements (in centimetres):

• Sepal length and width (the green outer part)
• Petal length and width (the colourful inner part)

Our task is to predict the species of a flower given these four measurements.

Loading and Exploring the Data

We start by loading the data from a Parquet file using the DataFrame library:

One immediate advantage of Haskell: the type system tells us we have no missing values. If any column contained nulls we’d see a Maybe prefix in front of that type. At least in this case, we don’t have to worry about deciding between imputing data and dropping it.

Exploring our data

Checking for balance

Before building any model, we need to know if our dataset is balanced. A severely imbalanced dataset (say, 140 Setosas and 5 each of the others) would make accurate prediction nearly impossible.

Great! A perfect balance with 50 flowers of each species. This means we don’t have to spend time on rebalancing tricks.

Some light analysis

Next, we compare feature magnitudes per class. We unintuitively do this in a stacked bar chart. But since our plot is interactive and allows us to exclude features at the click of a button we throw them all in there then click around to look at their individual differences.

Our charts broadly suggest that Virginica > Versicolor > Setosa for most features, with one fun exception: Setosa has the largest sepal width.

To get a sense of the spread of the features we can do box plots per variety and compare them:

Since we don’t have box plots implemented as web plots we fall back to granite’s ASCII plot.

For Setosa:

For Versicolor

For Virginica:

The box plots reveal that Setosas should be easy to identify—their petals are tiny with little variance. Versicolors and Virginicas overlap more, but we can still try to distinguish them using feature engineering.

Creating new features

From our exploration we have 5 promising feature candidates:

• sepal area
• sepal length to width ratio
• petal area
• petal length to width ratio
• sepal area / petal area

To derive these feature it’ll help to create typed references to the columns using template Haskell. This is the equivalent of doing a database migration against a schema definition. Except in this case we just freeze the types as they exist in the dataframe already.

:set -XTemplateHaskell
import qualified DataFrame.Functions as F

F.declareColumns df

import DataFrame ((|>))

engineered = df
               |> D.derive "sepal_area" (sepal_length * sepal_width)
               |> D.derive "sepal_ratio" (sepal_length / sepal_width)
               |> D.derive "petal_area" (petal_length * petal_width)
               |> D.derive "petal_ratio" (petal_length / petal_width)
               |> D.derive "area_ratio" ((sepal_length * sepal_width) / (petal_length * petal_width))

After repeating the previous diagramming (see the full notebook) we decide we can get by with three features:

• petal_area,
• sepal_area, and,
• area_ratio

Creating a custom data type for our label

Here’s where Haskell’s type system shines. Instead of treating species as magic integers, we define them explicitly:

data Iris
    = Setosa
    | Versicolor
    | Virginica
    deriving (Eq, Show, Read, Ord, Enum)

The deriving clause automatically generates useful functions:

• Eq: Allows us to compare Iris values for equality
• Show: Converts Iris to a String (e.g., “Setosa”)
• Read: Converts a String to Iris (e.g., “Setosa” → Setosa)
• Ord: Allows ordering/sorting
• Enum: Lets us convert to/from integers (Setosa=0, Versicolor=1, Virginica=2)

We convert our string labels to this type:

import qualified Data.Text as T

modellingDf = D.select ["petal_area", "sepal_area", "area_ratio", "variety"] engineered
withTypedLabel =
        modellingDf
            |> D.derive
                "variety"
                (F.lift (fromEnum . read @Iris . T.unpack) variety)

Preparing our data for ML

For our model to generalise, we split the data: 70% for training, 30% for testing. Our random split function takes a seed and probability p. It return a tuple where (p * 100) % of the vakues are in the first element and the rest are in the second.

The random seed (42) ensures reproducibility which is crucial for debugging.

import qualified DataFrame.Hasktorch as DHT
import qualified System.Random as SysRand
import Control.Exception (throw)

let (trainDf, testDf) = D.randomSplit (SysRand.mkStdGen 42) 0.7 withTypedLabel

let trainFeaturesTr =
        trainDf
            |> D.exclude ["variety"]
            |> DHT.toTensor
let testFeaturesTr =
        testDf
            |> D.exclude ["variety"]
            |> DHT.toTensor

let trainLabels = either throw id (D.columnAsIntVector "variety" trainDf)
let testLabels = either throw id (D.columnAsIntVector "variety" testDf)

Since we are predicting one of many classes, our target should be a one-hot vector:

• 0 (Setosa) → [1.0, 0.0, 0.0]
• 1 (Versicolor) → [0.0, 1.0, 0.0]
• 2 (Virginica) → [0.0, 0.0, 1.0]

import qualified Torch as HT

let trainLabelsTr = HT.toType HT.Float (HT.oneHot 3 (HT.asTensor trainLabels))
let testLabelsTr = HT.toType HT.Float (HT.oneHot 3 (HT.asTensor $ testLabels))

The function application can be read from right to left:

• Make the label into a tensor
• Create a one-hot vector with three entries.
• Then finally convert each int to a float.

That’s it. Our data is ready for machine learning. Now we can pass it to Torch and use it to train a small multi-layer perceptron.

Training the model

Neural Network Architecture

We define our Multi-Layer Perceptron (MLP) architecture in two parts.

First, a specification that describes the shape of our network. Second, the actual model with its layers. Each layer is a Linear transformation (like nn.Linear in PyTorch).

:set -XDeriveGeneric
import qualified Torch as HT

import GHC.Generics (Generic)

data MLPSpec = MLPSpec
    { inputFeatures :: Int   -- Number of input features (3 for our dataset)
    , hiddenFeatures :: Int  -- Number of neurons in hidden layer
    , outputFeatures :: Int  -- Number of output classes (3 species)
    }
    deriving (Show, Eq)

data MLP = MLP
    { l0 :: HT.Linear  -- Input → Hidden layer
    , l1 :: HT.Linear  -- Hidden → Output layer
    }
    deriving (Generic, Show)

Network Architecture Diagram:

Input Layer (3)  →  Hidden Layer (8)  →  Output Layer (3)
---------------     -----------------     ----------------
sepal_area          ReLU activation       Softmax
petal_area          (introduces           (produces
area_ratio          non-linearity)        probabilities)
                                          Setosa
                                          Versicolor
                                          Virginica

Making Our Model Trainable

We need to tell Hasktorch how to initialise our network with random weights. This is similar to defining __init__() in a PyTorch nn.Module:

:set -XRecordWildCards

instance HT.Parameterized MLP
instance HT.Randomizable MLPSpec MLP where
    sample MLPSpec{..} =
        MLP
            <$> HT.sample (HT.LinearSpec inputFeatures hiddenFeatures)
            <*> HT.sample (HT.LinearSpec hiddenFeatures outputFeatures)

The <$> and <*> operators are Haskell’s way of working with random initialisation. Think of this as: “Create an MLP by randomly sampling weights for both layers.”

Forward Pass

This function defines how data flows through the network. It’s equivalent to the forward() method in PyTorch. Read it from right to left (or bottom to top in the chain):

mlp :: MLP -> HT.Tensor -> HT.Tensor
mlp MLP{..} =
    HT.softmax (HT.Dim 1)       -- 4. Apply softmax (probabilities sum to 1)
        . HT.linear l1          -- 3. Apply second linear layer
        . HT.relu               -- 2. Apply ReLU activation
        . HT.linear l0          -- 1. Apply first linear layer

Read this right-to-left (or bottom-to-top). The . operator composes functions.

In Python/PyTorch, this would look like:

def forward(self, x):
    x = self.l0(x)
    x = F.relu(x)
    x = self.l1(x)
    x = F.softmax(x, dim=1)
    return x

The Haskell version is directly captures that we are composing a series of functions and using the output of one as the input to the other.

Training Loop

Training longer doesn’t always mean better. At some point, your model starts memorizing the training data instead of learning general patterns.

We detect this by watching the test set loss. As long as test loss keeps decreasing, the model is learning generalizable patterns. But if training loss keeps dropping while test loss starts climbing, we’re overfitting.

This is called early stopping, and it’s one of the simplest yet most effective regularisation techniques in machine learning.

The Training State

Our training loop needs to track more than just the current model. We maintain:

• Current model: Updated every iteration
• Best model: The weights that gave us the lowest test loss
• Best loss: What that test loss was
• Patience counter: How many checks since we last improved?

If we don’t improve for several consecutive checks (in our case 5 iterations), we give up and return the best model we found.

The loop structure might look unusual if you’re coming from PyTorch. Instead of a for loop with mutation, we use the very similar looking foldLoop which takes a number of iterations, an initial state and a function that evolves that state with each iteration. Each iteration receives the previous state and returns the new state. No hidden mutable variables, just explicit data flow.

This is eventually the sort of thing we should abstract behind a good machine learning library but for now it’s easy enough to implement.

import Control.Monad (when)

trainLoop ::
    Int ->                          -- Number of epochs
    (HT.Tensor, HT.Tensor) ->       -- Training features and labels
    (HT.Tensor, HT.Tensor) ->       -- Test features and labels
    MLP ->                          -- Initial model
    IO MLP                          -- Returns trained model
trainLoop n trainingData testData initialModel = do
    let initialState = makeInitialState initialModel
    (_, bestModel, _, _) <- HT.foldLoop initialState n (trainingStep trainingData testData)
    pure bestModel
  where
    -- Configuration constants
    patience = 5
    checkInterval = 500
    learningRate = 1e-2
    initialBestLoss = read @Float "Infinity"
    
    -- Create initial training state
    makeInitialState model = (model, model, initialBestLoss, 0)
    
    -- Main training step for each iteration
    trainingStep (features, labels) (testFeatures, testLabels) (model, bestModel, bestLoss, counter) i
        | shouldStopEarly counter = pure (model, bestModel, bestLoss, counter)
        | otherwise = do
            -- Train and evaluate
            (updatedModel, trainLoss) <- performTrainingStep model features labels
            testLoss <- computeTestLoss updatedModel testFeatures testLabels
            
            -- Log progress periodically
            when (shouldLogProgress i) $
                logTrainingProgress i trainLoss testLoss
            
            -- Update best model tracking
            let (newBest, newBestLoss, newCounter) = 
                    updateBestModelTracking updatedModel testLoss bestModel bestLoss counter
            
            pure (updatedModel, newBest, newBestLoss, newCounter)
    
    -- Early stopping condition
    shouldStopEarly counter = counter >= patience
    
    -- Logging condition
    shouldLogProgress i = i `mod` checkInterval == 0
    
    -- Perform one training iteration
    performTrainingStep model features labels = do
        let predictions = mlp model features
            loss = HT.binaryCrossEntropyLoss' labels predictions
        (updatedModel, _) <- HT.runStep model HT.GD loss learningRate
        pure (updatedModel, HT.asValue loss :: Float)
    
    -- Evaluate model on test set
    computeTestLoss model features labels =
        let predictions = mlp model features
            loss = HT.binaryCrossEntropyLoss' labels predictions
        in pure (HT.asValue loss :: Float)
    
    -- Print training progress
    logTrainingProgress iteration trainLoss testLoss =
        putStrLn $ "Iteration: " ++ show iteration
                ++ " | Training Set Loss: " ++ show trainLoss
                ++ " | Test Set Loss: " ++ show testLoss
    
    -- Update best model if test loss improved, otherwise increment patience counter
    updateBestModelTracking currentModel currentLoss bestModel bestLoss counter =
        if currentLoss < bestLoss
            then (currentModel, currentLoss, 0)  -- New best: reset counter
            else (bestModel, bestLoss, counter + 1)  -- No improvement: increment counter

Running the training loop

We bring everything together by running our training loop with a model intialized from our MLPSpec.

initialModel <- HT.sample $ MLPSpec 3 8 3

trainedModel <-
    trainLoop
        10000
        (trainFeaturesTr, trainLabelsTr)
        (testFeaturesTr, testLabelsTr)
        initialModel

putStrLn "Your model weights are given as follows: "
print trainedModel

Evaluating the model

The notebook creates some helps to evaluate to the model. We’ll skip these details since most of the code is vanilla Haskell.

And for our results. Training metrics:

------------------------------------------
index |  variety   | precision |  recall  
------|------------|-----------|----------
 Int  |    Iris    |   Float   |   Float  
------|------------|-----------|----------
0     | Setosa     | 1.0       | 1.0      
1     | Versicolor | 0.9677419 | 0.9375   
2     | Virginica  | 0.9459459 | 0.9722222

Test metrics:

------------------------------------------
index |  variety   | precision |  recall  
------|------------|-----------|----------
 Int  |    Iris    |   Float   |   Float  
------|------------|-----------|----------
0     | Setosa     | 1.0       | 0.9285714
1     | Versicolor | 0.8888889 | 0.8888889
2     | Virginica  | 0.8666667 | 0.9285714

Conclusion

This post demonstrates building a complete multiclass classification system in Haskell from scratch. We tackle the classic Iris dataset. The journey covers data loading, exploratory analysis, feature engineering, typed data modelling, network architecture design, training with early stopping, and evaluation.

The tools for a robust data science journey are scattered in the Haskell ecosystem. The hope is that we can unify them and along the way create ergonomic and safe APIs for data science.

Wanna join in this work?

Come over to the dataHaskell Discord to help shape the future of data science and Haskell?

The downloaded ipynb for this post can be found here.

This post kicks off a hands-on series about program synthesis—the art of teaching machines how to generate code. We’ll build a tiny, FlashFill-style synthesiser that learns to turn strings like “Joshua Nkomo” into “J. Nkomo” from input/output pairs. We’ll see how to define a tiny string-manipulation language, write an interpreter, and search the space of programs to find one that solves our toy problem.

The accompanying code for this post is on Github and runs in under a second for most examples.

Why synthesis matters

System 1 vs System 2 thinking

The book “Thinking Fast and Slow” by Daniel Kahneman differentiates between two kinds of thinking: system 1 vs system 2 thinking.

System 1 thinking is characterised by fast, intuitive judgments. It’s the thinking that you do when you’re not fully aware that you’re thinking. It is built up through repetition and experience. How we express language is an example of System 1 at work. In most situations, we don’t consciously think of what words to say. They stream out of us, dressing the thoughts we intend to communicate. System 1 thinking is the realm of the familiar. It’s not always right, but it gets you into a comfortable ballpark.

System 2 thinking, on the other hand, is slow and deliberate. It’s speaking in front of a crowd for the first time. It’s taking an exam for a class. It’s how we think when a context is outside (or at the bounds) of what we are used to. System 2 solves hard problems, but is time-consuming and taxing.

Human cognition is a constant dance between the two. Talking to a loved one is easy until the conversation topic is something sensitive. Writing a proof is hard until you discover the “trick” that could lead you to an answer. The two systems also interact in very strange ways. Ever had a eureka moment in the shower? Ever gotten a joke while driving, days after it was told? Problem-solving often means alternating between these two systems of thought.

How machines “think”

An analogous dichotomy can be applied to artificial intelligence.

Machine learning (broadly defined as using statistics to learn patterns from data) is system 1 thinking. It is the statistical “intuition” gained from gorging on mounds of data. It’s easy to invalidate a lot of statistical approaches as memorisation (a lot of people will say LLMs are just a fancy autocomplete), but building the right statistical models is as difficult as building the right intuitions to solve problems. There has been a lot of ingenuity in this space over the last few decades.

But that’s not the entire story.

Until the advent of deep learning in the early 2010s, artificial intelligence was dominated by systems that tried to recreate system 2 thinking. These are “slow” algorithms that take time to discover the “right” answer. Program synthesis is one such approach. It tries to search through the space of instructions until it discovers one that solves a specified problem.

A growing number of researchers believe that getting to artificial general intelligence involves marrying the two approaches (what’s now referred to as neuro-symbolic AI). There have been some promising achievements in the field recently:

DeepMind’s AlphaGeometry combines a learned model with a symbolic deduction engine and massive synthetic training data to tackle Olympiad-style geometry: in a Nature-reported benchmark of 30 problems drawn from IMO 2000–2022, it solved 25 within the time limit, approaching a human gold-medalist’s performance. Later, an upgraded AlphaGeometry 2 paired with AlphaProof (a reinforcement-learning formal prover) solved 4 of 6 problems from IMO 2024—silver-medal level. If you want a quick, approachable walkthrough of how this works, Aleph0 has a really good video on this.

On the “reasoning LLM” front, OpenAI’s o3 pushed ARC-AGI scores sharply upward. François Chollet characterizes o3’s core as a form of deep learning guided program search: at test time the system explores and evaluates many candidate chains-of-thought—effectively short “programs” written in natural language—guided by a learned prior, and selects the best-scoring plan. Regardless of the exact internals, the result is stronger adaptation to novel tasks than earlier GPT-style models.

Stepping back, these examples point to a common pattern: use learning to propose promising moves, and use search/logic to compose, check, and refine them. That division of labor (neural intuition plus symbolic rigor) is increasingly the template for systems that need both creativity and reliability.

So, as we stand on the cusp of a future where synthesis looks as though it will make a resurgence in some incarnation, I thought it apt to write a series of practical blog posts on what it is.

What is program synthesis

We’ll work from a definition offered in Armando Solar-Lezama’s MIT class on the subject.

Program Synthesis correspond to a class of techniques that are able to generate a program from a collection of artifacts that establish semantic and syntactic requirements for the generated code.

Program synthesis lives at the intersection of what you want (the specification) and how you look for it (the search). Syntheis searches through the space of candidate programs using a given criteria.

Searching and Specification

Specifying the problem

A specification is any collection of clues that narrows down the set of candidate programs. Those clues can be semantic (what the output must mean) or syntactic (what the code must look like).

Synthesis problems are usually specified through:

• Input–Output Examples (I/O): The friendliest form. “When input is Joshua Nkomo, output must be J. Nkomo.”
• Constraints / Properties (pre/postconditions): Rules like “output must contain a full-stop,” “output length ≤ max(len(input), len(target)),” or “index must be within bounds.” These semantic checks prune huge regions of bad candidates even before running full tests.
• Types and Grammars (the DSL itself): By choosing a small, typed vocabulary (our string blocks), you ban ill-formed programs outright. This is a syntactic specification baked into the language design.
• Sketches (templates with holes): “I’m pretty sure it’s Concat(Upper(Head(x)), ???)—fill in the ???.” You constrain structure and let the tool finish the details.
• Negative examples: “For input Steve, don’t output S. with a trailing space.” Prevents sneaky overfits.
• Resource hints / preferences (ranking): “Prefer shorter programs,” “penalize Substring with two dynamic indices,” etc. These are soft constraints used to rank ties.

In our mini Flash-Fill program, the specification will be I/O examples.

Searching the space

At its core, synthesis is a search problem. Several families of search exist:

1. Enumerative search
2. Constraint-guided search
3. Heuristic / ML-guided search

We’ll focus on the latter two in subsequent posts. For this post we’ll use a simple, yet powerful, type of enumerative search.

Building Our Synthesiser

Follow along as you read the code on Github

Starting from the bottom

Here’s the key insight: if we have a language of simple operations (head, tail, concat, etc.), we can systematically combine them to build complex programs. It’s like having LEGO blocks - start with basic pieces and build increasingly complex structures until one solves the problem.

We’ll write a synthesiser that generates programs that can solve problems of the following nature:

The building blocks

We usually start by defining what the “legal moves” of our problem are. The “program” part of program synthesis is typically a small, expressive set of instructions that can solve our specific problem. We call this set of instructions a domain specific language (DSL).

It is easy to view DSLs as a limitation. Why define a small language instead of generating any arbitrary code in a language like Python or Java? Doesn’t a DSL mean that our technique inherently lacks power and expressivity?

While a DSL can be viewed as a constraint, its contribution is far more profound than a simple restriction. The design of the DSL is deeply intertwined with the design of the search algorithm. A well-designed DSL exposes a set of operators and a program structure that are amenable to efficient search and, crucially, aggressive pruning of the search space. DSLs are how we embed domain knowledge into program synthesis.

The development of FlashFill provides a canonical example of this principle. The synthesiser’s DSL was not chosen arbitrarily; it was carefully engineered to enable an efficient top-down, divide-and-conquer search algorithm.

If a DSL is too small, we risk not being able to express the target transformation. If it’s too big, we drown in candidates and overfit to examples. We want a sweet spot: just expressive enough to cover the problem class, yet constrained enough that searches can complete in our lifetime.

After eyeballing our problem (string manipulation) we decide our language should support these operations:

• Head — keep the first character (“Joshua” → “J”)
• Tail — drop the first character (“ Nkomo” → “Nkomo” after trimming a leading space)
• Lower / Upper — change case
• Concat — glue two strings together
• Substring start end x — slice a range from x
• Find needle haystack — index of needle inside haystack
• Constants like “.”, “ “
• A placeholder input value

That’s enough to express lots of “data cleaning” tricks.

Given our mini-language, we can express our intended program as follows:

firstInitial = Upper(Head(x))

-- Find the occurrence of the space, then drop it.
lastName = Tail( Substring( Find(" ", x), End, x ) )

Concat (
 Concat ( firstInitial, ". "),
 lastName
)

-- Or written densely.
Concat(
 Concat( Upper(Head(x)), ". " ),
 Tail( Substring( Find(" ", x), End, x ) )
)

Our challenge is to see if we can come up with a similar program using synthesis.

First, we have to define our language programmatically:

data Program a where
 Concat :: Program String -> Program String -> Program String
 Substring :: Program Int -> Program Int -> Program String -> Program String
 Tail :: Program String -> Program String
 Head :: Program String -> Program String
 SValue :: String -> Program String
 Variable :: Program String
 Find :: Program String -> Program String -> Program Int
 Start :: Program Int
 End :: Program Int

Program a means “a little recipe (an expression tree) that, when we run it, produces a value of type a.” In our code, there are only two kinds of results: Program String and Program Int.

The :: Program String or :: Program Int parts tell us what kind of result each operation produces:

• Program String = produces text
• Program Int = produces a number (usually a position)

The arrows (->) show what inputs each operation needs. For example:

• Tail :: Program String -> Program String means “Tail needs a string-producing program as input and produces a string-producing program as output”

This structure is compositional. We can define longer programs in terms of these basic elements as we did in the solution above. The synthesis community typically uses functional programming languages because they express these relationships more clearly.

The interpreter for our DSL can also be defined recursively:

-- | This is the leaf of our tree - string wrapped as a program
-- or rather a program that returns a string constant.
interpret (SValue v) = v
-- | Our string operations
interpret (Tail v) = drop 1 (interpret v)
interpret (Head v) = take 1 (interpret v)
interpret (Concat l r) = interpret l ++ interpret r
interpret (Substring start end g) = ...

Think of each Program as a tree node: SValue "foo" is a leaf; Tail v, Head v have one child; Concat l r has two children; Substring start end g has three. interpret walks that tree. Each line says, “if the node looks like this shape, do this.” Haskell lets us write one equation per shape of input; the one that matches is used.

Now we can define our also search strategy programatically. To recap:

1. Start small. Begin with the identity function and the input x. Optionally, seed a few helpful constants (like “ “ or “.”) and any characters shared across all inputs/outputs.
2. Grow candidates. Repeatedly build bigger programs by:
   • Wrapping existing ones with Head, Tail, etc.
   • Concatenating any two existing ones.
   • Making substrings using Find, Start, and End.
3. Interpret and test. Run each candidate on the provided examples; keep only those that don’t explode in size and (eventually) those that match all examples.
4. Iterate by depth. Increase the “size limit” and repeat until we succeed or give up.

search ::
 -- | Examples
 [(String, String)] ->
 -- | Variables of interest.
 [Program String] ->
 -- | Existing Programs
 [Program String -> Program String] ->
 -- | Search depth
    Int ->
 -- | Maybe is an optional type meaning we might end with no program found.
    Maybe (Program String -> Program String)
search examples variables programs d =
    if d == 0 then Nothing  --| Give up if we've gone deep enough
    else
        case findFirst ps of
 Just p -> Just p
 Nothing -> search examples variables (generatePrograms variables programs) (d - 1)
  where
 findFirst [] = Nothing
 findFirst (p : ps')
 | satisfiesExamples examples p = Just p
 | otherwise = findFirst ps'

The where section defines some variable bindings, and he core of the logic is in the if statement.

The entire system is recursive:

• Generation 0 produces Generation 1
• Generation 1 produces Generation 2
• Each generation builds on all previous work

Notice, specifically, that the input to the new search is all the generated programs from the current generation, and the other inputs (except for the depth parameter) are left unchanged.

search examples variables (generatePrograms variables programs) (d - 1)

Generate Programs Systematically

The generatePrograms function in our code employs four distinct strategies simultaneously, each exploring a different dimension of the program space:

Strategy 1: Simple Transformations

[ comp transform p | p <- existingPrograms, transform <- [Tail, Head, Lower, Upper] ]

This strategy takes every existing program and applies simple one-step transformations. If we have a program p that produces “HELLO”, this generates:

• Tail(p) → “ELLO” (remove first character)
• Head(p) → “H” (keep only first character)

Strategy 2: Concatenation Combinations

[ \v -> Concat (p v) (q v) | p <- existingPrograms, q <- existingPrograms ]

This creates programs that combine the outputs of two existing programs. If we have:

• Program p: extracts first word
• Program q: extracts last word

This generates Concat(p, q): first word + last word

Example Evolution:

• Generation 1: p = identity (returns input as-is)
• Generation 2: Concat(p, p) (doubles the input: “hi” → “hihi”)
• Generation 3: Concat(Concat(p, p), p) (triples: “hi” → “hihihi”)

Combinatorial Explosion: With just 10 existing programs, this strategy alone generates 10 × 10 = 100 new programs! This is something we must contain to make this approach useful.

Strategy 3 & 4: Variable-Based Transformations

These are the most sophisticated strategies. They use discovered “variables” (important strings found in the examples) to create complex transformations.

[ comp t p | v <- vars, p <- existingPrograms, t <- transformsWithVar v ]
[ comp p t | v <- vars, p <- existingPrograms, t <- transformsWithVar v ]

Let’s say the system discovered “@” as an important variable. It generates programs like:

Pre-processing with variable (Strategy 3):

1. Concat("@", existing_program) - Prepend @ to result
2. Substring(0, Find("@", input), existing_program) - Everything before @
3. Substring(Find("@", input), End, existing_program) - Everything after @

Post-processing with variable (Strategy 4):

1. existing_program then Concat with “@”
2. existing_program then extract until “@”
3. existing_program then extract after “@”

A Concrete Generation Example

Let’s trace through generating a program for email domain extraction:

Initial State (Generation 0)

• Programs: [identity] (just returns input)
• Variables: ["@", "."] (discovered from examples)

Generation 1

The system generates programs like:

• Tail(identity) → removes first char
• Head(identity) → keeps first char
• Concat(identity, identity) → doubles the input
• Concat("@", identity) → prepends @
• Concat(identity, "@") → appends @
• Substring(Find("@", identity), End, identity) → everything after @

If the solution required more complexity, Generation 2 would create even more programs:

• Tail(Substring(Find("@", identity), End, identity))
• Concat("domain: ", Substring(Find("@", identity), End, identity))
• Tail(Substring(Find("@", identity), End, identity))

The final function looks like this:

-- | Generate the next “ring” of candidate programs from:
--   * vars: string-valued building blocks discovered from examples (e.g., " ", ".", "@", etc.)
--   * existingPrograms: functions that turn an input Program String into a Program String
--   The result is a list of new transformers, ordered (smallest AST first) by a simple size metric.
generatePrograms
 :: [Program String]                      -- ^ vars (constants/placeholders from examples)
 -> [Program String -> Program String]    -- ^ existing candidate transformers
 -> [Program String -> Program String]    -- ^ newly generated transformers (sorted by size)
generatePrograms vars existingPrograms =
  -- Sort candidates so we explore simpler/smaller programs first.
  -- We estimate size by applying each transformer to a dummy leaf (SValue "")
  -- and measuring the generated AST with gSize.
  L.sortBy
 (\p q -> gSize (p (SValue "")) `compare` gSize (q (SValue "")))

    -- 1) Unary wrappers: for every existing transformer p,
    --    wrap it once with a simple string op (Tail/Head/Lower/Upper).
    --    `comp transform p` means “apply p, then apply transform”.
 ( [ comp transform p
 | p <- existingPrograms
 , transform <- [Tail, Head, Lower, Upper]
 ]

    -- 2) All ordered concatenations of existing transformers:
    --    build a function that, given an input v, concatenates (p v) and (q v).
    --    This grows breadth by combining previously discovered behaviours.
 ++ [ \v -> Concat (p v) (q v)
 | p <- existingPrograms
 , q <- existingPrograms
 ]

    -- 3) Variable-aware transforms (pre-compose):
    --    for each discovered var v (a Program String), and each existing p,
    --    create transforms t that *mention* v (prepend/append/slice around v),
    --    then do `t . p` via `comp t p`.
 ++ [ comp t p
 | v <- vars
 , p <- existingPrograms
 , t <- transformsWithVar v
 ]

    -- 4) Variable-aware transforms (post-compose):
    --    symmetric to (3), but apply p first and then t (`p . t` via `comp p t`),
    --    so we explore both “before p” and “after p” placements of the same motif.
 ++ [ comp p t
 | v <- vars
 , p <- existingPrograms
 , t <- transformsWithVar v
 ]
 )
  where
 -- | Small library of transforms parameterised by a discovered variable `v`.
 --   Each lambda expects a string-producing Program `v'` (usually the current input),
 --   and returns a new string-producing Program that uses `v` in a useful way.
 transformsWithVar v =
 [ \v' -> Concat v v'                           -- prepend v
 , \v' -> Concat v' v                           -- append v
 , \v' -> Substring Start (Find v v') v'        -- prefix of v' up to first v
 , \v' -> Substring Start (Find v' v) v         -- prefix of v up to first v'
 , \v' -> Substring (Find v v') End v'          -- suffix of v' after first v
 , \v' -> Substring (Find v' v) End v           -- suffix of v after first v'
 ]

-- | Helper intuition (defined elsewhere in our code):
-- comp :: (Program String -> Program String)
--      -> (Program String -> Program String)
--      -> (Program String -> Program String)
-- comp = (.)
--
-- gSize :: Program String -> Int
-- gSize = ... -- counts AST nodes to bias search toward simpler programs

Our approach is like evolution - each generation inherits and combines traits from its ancestors, occasionally producing exactly what we need. Given enough time and memory, it will find all relevant programs expressible in the language.

The key caveat here, though, is “enough time and memory.” We’d like to find these programs in our lifetime. How do we tame the growth of our generation function?

Our program, as it exists, runs fine for small transformations, e.g. [("Joshua", "J"), ("Steve", "S")], but for our target problem, my computer runs out of memory and kills the program (even after running for close to an hour).

We need to be smarter about how we search the space/

The Exponential Growth Problem

Exponential growth is the supervillain of the synthesis story. Looking at the code above, we can see why! We generate A LOT of programs at each step.

Let’s calculate the growth rate:

Generation 0: 1 program (identity) Generation 1: ~30 programs (depending on variables discovered) Generation 2: ~900 programs (30 existing × 30 combinations) Generation 3: ~27,000 programs Generation 4: ~810,000 programs

The sophistication of an enumerative synthesiser is largely measured by its ability to intelligently prune the search space. Pruning strategies exist on a spectrum, from reactive deduplication of redundant programs to proactive elimination of entire branches of the search tree. L

Let’s rewrite our search to include some extra pruning and deduplication steps as examples.

search examples variables programs d
 | d == 0 = Nothing
 | otherwise =
        case findFirst ps of
 Just p -> Just p
 Nothing -> search examples variables (generatePrograms variables ps) (d - 1)
  where
 inputs = map fst examples  -- | separate the inputs from the outputs.
 ps = prune examples $ deduplicate inputs programs  -- | Now define a new set of deduplicated, pruned programs.
 findFirst [] = Nothing
 findFirst (p : ps')
 | satisfiesExamples examples p = Just p
 | otherwise = findFirst ps'