Swapping a data structure at runtime with coheriq

coheriq lets a library mark certain functions and data structures as candidates for acceleration, and lets a separate provider swap in alternative implementations at runtime — without the calling code changing at all.

There are three parties in this example, as is typical:

1. The domain — a library (coheriq_ragged_domain) that marks candidates and ships working pure-Python defaults.

2. The engine — a separate package (coheriq_ragged_engine) that provides a faster, compiled implementation (C++ via nanobind, in this case), discovered through a Python entry point.

3. User code — activates an engine once, then calls the library exactly as before.

This notebook demonstrates the mechanism: a container class and the functions that operate on it are swapped, and the calling code is unchanged.

It is written to run either way — against the pure-Python default or the compiled engine — without editing any calling code. You choose the backend by activating the engine (see below) or by setting the RAGGED_ENGINE=fused environment variable before launching. Run it both ways and compare the timing printed near the end: same code, different backend.

The source code to the domain and engine are available in the example/ directory of the coheriq repository.

The shape: ragged data

Our computation is a per-row softmax over a batch of rows that have different lengths — a ragged batch. Ragged data is everywhere: variable-length sequences, a different number of candidates per item, the neighbors of each node in a graph.

This is exactly the shape that a rectangular array library like NumPy handles awkwardly. You have two options, both unpleasant:

• A Python list of arrays — back to per-row Python overhead, which defeats the vectorization that made NumPy worth reaching for.

• Pad every row to the longest and carry a mask — memory and compute become O(rows × max_len) instead of O(total elements), which is wasteful when row lengths vary a lot.

This is why libraries grew dedicated ragged/segmented machinery — tf.RaggedTensor, torch_scatter, JAX’s segment_* operations. The shape is real and general. coheriq gives your own library a clean way to swap in such a backend for its own data type, without your users rewriting anything.

Using the library (pure Python, no engine yet)

A user installs coheriq_ragged_domain and imports its container and functions.

import inspect

from coheriq_ragged_domain import RaggedBatch, segmented_softmax, segmented_topk_softmax

What the default looks like

coheriq resolves each candidate’s implementation on its first call and then freezes it for the rest of the process. So we must choose the backend before calling anything — which is why this notebook decides up front and then runs one way through. (If you call a candidate first and then try to enable an engine, coheriq deliberately raises to tell you it is too late.)

Here is the default implementation, shown by reading it (this always shows the domain’s own source, whichever backend is active). The container stores all rows in one flat buffer plus an offsets array (CSR-style), and segmented_softmax does the obvious per-row work — find the max, exponentiate, sum, divide — materializing a Python list for each row:

print(inspect.getsource(segmented_softmax))

@_domain.acceleration_candidate
def segmented_softmax(batch: RaggedBatch) -> RaggedBatch:
    """Return a new batch holding the per-row softmax of ``batch``.

    The default implementation is the obvious, readable version: for each row it
    makes several passes -- find the max (for numerical stability), subtract and
    exponentiate, sum, then divide.  Each row is materialized as a Python list
    along the way.
    """
    result_rows: list[list[float]] = []
    for row in batch:
        if not row:
            result_rows.append([])
            continue
        m = max(row)
        exps = [math.exp(x - m) for x in row]
        total = sum(exps)
        result_rows.append([e / total for e in exps])
    return RaggedBatch(result_rows)

Activating the engine — two ways, without importing it

A separate package, coheriq_ragged_engine, is installed. It declares itself under the entry-point group coheriq.engines.ragged with the name "fused". Its implementation is compiled C++ (via nanobind) — but you never import it.

There are two equivalent ways to activate it, and neither imports the engine package:

1. Explicitly, by name — uncomment the line in the next cell: coheriq.enable_engine("ragged", "fused").

2. Via the environment — set RAGGED_ENGINE=fused before launching; coheriq activates it automatically on first use.

Either way, coheriq looks up the entry point, imports the compiled engine for us, and enables it. Leave the line commented and the variable unset to run against the pure-Python default instead. Notice there is no ``import coheriq_ragged_engine`` anywhere in this notebook.

# Uncomment to activate the compiled engine explicitly (or set RAGGED_ENGINE=fused):
# import coheriq
# coheriq.enable_engine("ragged", "fused")

The swap: same calling code, whichever backend is active

We build a batch with wildly varying row lengths (including an empty row). The calling code is identical no matter which backend is active — the object we get back reports its backend via .backend, so we can see which implementation is live. If you activated the engine, it says fused; otherwise default.

batch = RaggedBatch([[2.0, 1.0, 0.1], [1.0], [3.0, 3.0, 0.0, 0.0], []])

print(repr(batch))
print("backend:   ", batch.backend)
print("len:       ", len(batch))
print("row_lengths:", batch.row_lengths())
print("batch[0] (a copy):", batch[0])

RaggedBatch(4 rows, default backend)
backend:    default
len:        4
row_lengths: [3, 1, 4, 0]
batch[0] (a copy): [2.0, 1.0, 0.1]

The operation is swapped too. segmented_softmax runs whichever backend is active, and the result it hands back reports the same backend — so a result from the compiled engine is itself a fused batch.

probs = segmented_softmax(batch)

print("result backend:", probs.backend)
for i in range(len(probs)):
    row = probs[i]
    print(f"row {i}: sum={sum(row):.6f}  {[round(x, 4) for x in row]}")

result backend: default
row 0: sum=1.000000  [0.659, 0.2424, 0.0986]
row 1: sum=1.000000  [1.0]
row 2: sum=1.000000  [0.4763, 0.4763, 0.0237, 0.0237]
row 3: sum=0.000000  []

Correctness by invariant

Different engines need not produce bit-identical output (a compiled backend might sum in a different order). So we check invariants rather than equality: the batch shape is preserved and every non-empty row is a valid probability distribution.

def check_softmax_invariants(inp, out, tol=1e-9):
    assert len(out) == len(inp)
    assert out.row_lengths() == inp.row_lengths()
    for i in range(len(out)):
        row = out[i]
        if row:
            assert abs(sum(row) - 1.0) < tol, (i, sum(row))
        assert all(-tol <= p <= 1.0 + tol for p in row)
    return True


assert check_softmax_invariants(batch, segmented_softmax(batch))
print("invariants hold; active backend is", probs.backend)

invariants hold; active backend is default

Where the compiled engine wins

When the fused engine is active, segmented_softmax runs a compiled C++ implementation (via nanobind) instead of the Python default. It exploits exactly the structure this data shape affords:

• A fused pass over contiguous storage — the flat values buffer is walked directly; no per-row temporary Python list is allocated.

• No Python object overhead — no boxed floats, tuples, or dynamic dispatch in the inner loop.

• Room for SIMD on the exponential and the arithmetic.

• Parallelism across rows — the rows are independent; the computation is embarrassingly parallel.

From the caller’s point of view nothing changes: the same segmented_softmax(batch) call dispatches to it. The benchmark near the end lets you measure the difference by running this notebook once with the default and once with the engine active.

Elaboration: top-k softmax

The same swap pattern applies to other operations on the container. A common variant keeps only the k largest entries per row and softmaxes over just those. The default sorts each row (O(n log n)); the compiled engine selects the top k with std::nth_element (O(n)) — an algorithmic win, not just a constant factor. Whichever backend is active, the call site is the same:

topk = segmented_topk_softmax(batch, k=2)
for i, row in enumerate(topk):
    print(f"row {i}: {[(idx, round(p, 4)) for idx, p in row]}")

row 0: [(0, 0.7311), (1, 0.2689)]
row 1: [(0, 1.0)]
row 2: [(0, 0.5), (1, 0.5)]
row 3: []

Benchmark: same code, two backends

The honest way to compare is to run this exact notebook twice — once with the default backend and once with the engine active (uncomment the activation cell, or set RAGGED_ENGINE=fused) — and compare the number printed below. The cell only prints its timing; it makes no assertion about it (hardware varies, and a compiled backend only pulls ahead once the rows are large enough to outweigh the Python call overhead), so it never fails a test run either way.

We build a large ragged batch with a fixed seed so both runs use identical input, then time segmented_softmax.

import random
import time

rng = random.Random(20260710)
big_rows = [[rng.gauss(0.0, 1.0) for _ in range(rng.randint(1, 128))] for _ in range(20_000)]
big = RaggedBatch(big_rows)

start = time.perf_counter()
result = segmented_softmax(big)
elapsed = time.perf_counter() - start

# Correctness is checked regardless of backend; timing is only reported.
assert check_softmax_invariants(big, result)
print(f"backend={big.backend}  rows={len(big)}  elapsed={elapsed * 1e3:.1f} ms")

backend=default  rows=20000  elapsed=182.8 ms

Recap

• A library (domain) marked a container and some functions as acceleration candidates and shipped working pure-Python defaults.

• A separate engine package provided a compiled C++ implementation and was discovered via a Python entry point — we activated it by name (or via an environment variable) and never imported it.

• The user code was unchanged across the swap; both the RaggedBatch container and the functions operating on it were replaced underneath it.

• We proved the swap with an observable marker (.backend), checked correctness by invariant, and measured the difference by running the same notebook against each backend.