The answer is never the one in the file extension. Every speed-up below is a rewrite that computes exactly the same thing as the version before it — same language, same file, same result — which makes it a compiler pass, performed by hand, in the source layer, because that was the only layer I was allowed to stand on. So the question is really: who should have written it, me or the compiler?

Put plainly: the different optimisations I show below within a single language are not performance tips. They are the demonstration that the abstraction is leaking: you are standing in for the compiler, transforming one function into another that computes the same thing and runs at a different speed, because the compiler sees the two differently.

Seen that way, Julia and JAX are powerful for opposite reasons. Julia’s metaprogramming lets you program the compiler, so instead of transforming the code by hand you build the abstraction that lets the compiler do the transforming. JAX goes the other way: it is restrictive enough that hand-transforming is not on the table, so you write it idiomatically and the compiler does the rest. Neither is free — Julia still needs you to find the transformation first, and JAX’s restriction only pays inside the domain it covers — but Numba does neither, which is why it is where the leak shows most plainly.

Numba, JAX and Julia are three different answers to that — not three speeds. Each comes with its own tower of abstractions, its own restricted dialect, and its own idea of what you are allowed to say, and the optimisation strategy arrives bundled with the rest rather than being something you pick separately.

So: one long section of concrete examples, three JIT compilers taken one at a time; then a step back into the more abstract way of talking about them; then, at the end, the surprising bit.

CPython’s own is the only one on that list you have to opt into — PEP 744, still experimental and off by default. Everything else is running whether you picked it or not.

Here is a very simple example: a one-line Python function doing some linear algebra, and the change you have to make to get it just-in-time compiled. It is not much of a change — you add a decorator. The signature I have written into it is unnecessary in this case; a bare @jit works here. I like being specific.

The compiled version is about 2.4× faster than the pure NumPy one, and the reason is where the abstraction falls. In Python every operation in that expression goes through an operator, and every operator creates a new object. A @ B multiplies the two and hands you back a new array. α * that gives you another new array. β * C gives you a third. The addition itself gives you a fourth. You need one thing at the end, and you have gone through a lot of memory allocation to get there. Put the decorator on and Numba is able to see through most of it and cut the allocations down.

The one it cannot cut is the matrix multiplication. Numba still does two allocations here, and in principle you only need one. A @ B gets dispatched to a linear algebra library, which allocates its own output, and Numba cannot see past that operator to avoid it. (The notebook checks the 4 against 2 in three independent ways: arrayexpr nodes in the typed IR, tracemalloc peak bytes, and an ndarray subclass counting __array_finalize__.)

The other thing worth highlighting is that what happens behind the scenes is very complicated, even though all we did was add one decorator. What Numba sees is not your source code. CPython has already compiled the function to bytecode, so Numba sees the bytecode first, and it has to be able to understand that bytecode — in some cases manipulating it — to work out what you wrote. Then it keeps lowering: to Numba’s own IR, then to a typed Numba IR, then to LLVM IR, then to assembly. It exposes every rung on the way down, through .inspect_types(), .inspect_llvm() and .inspect_asm(), which is what the notebook walks through.

This one is more complicated, and it is built to show two things: that in Numba you can write loops and still be quite fast, and that here writing the loop is faster — which is counter-intuitive to what we are trained to believe when we write array code. In NumPy you optimise by vectorising, by thinking about broadcasting and not using loops. The kernel is a Mexican hat over a grid.

Adding @jit to the same source is only slightly faster, and none of these numbers move very far, because the ceiling here is memory bandwidth rather than arithmetic. The column to read is peak, not the time: 4 → 2 → 1 is the structural result, and the two ends of the table follow it. The middle row does not — f_numba_inlined halves the peak and is still slower than the plain @jit — because peak and allocation count answer different questions. It trades the 127 MB r2 for two kilobyte-sized temporaries, so it holds half the memory while allocating three arrays instead of two. Only f_numba_loops wins on both counts.

f_numba_loops is the most optimised version, and what it is doing is a lot of low-level optimisation. First, the loops are explicit. Rather than implicitly assuming everything is an array operation and letting broadcasting sort it out — both inputs are 1D, the result is 2D, so they have to be broadcast along different axes — you write down the for i and the for j yourself. Then there is hoisting: you spot a computation that would otherwise be done over and over again and you lift it out of the loop. The part involving x has nothing to do with the inner loop, so it goes above it, and that is one less computation per element. This is a very typical example of how you would optimise a function in C. You could write exactly the same algorithm in C and it would look very similar.

Now, why does the naive @jit version still allocate twice? Both of those lines — r2 = ... and the return — are individually understood by Numba as array expressions. What it cannot see past is fusing the two together. r2 gets its own memory allocated, and then that r2 is used in the second line for the next computation. Ideally you want the intermediate gone, and in the loop version it is; but the compiler, looking at the first function, does not see past the array expression. So r2 becomes a 3600×4400 array, and that is a 127 MB round trip through memory that buys you nothing.

Two questions came from the room here.

“Does that mean it can see into the exponential function?” It stays a call — it is not going to be expanded and lowered and optimised further. But that is not what this is about. The point is whether it can see that all of this is an elementwise operation. Expand that second line and it is for i, for j, repeating the same expression and putting out an entry, every element independent of every other. That is what a compiler ought to be able to see and act on, and if it does, it knows it never needed to create the intermediate r2 array at all. Written the way it is written, the compiler cannot understand that. Could you write a compiler that sees past it? Yes, actually — there is an example a couple of sections down. Just not inside the abstraction Numba has built.

“Is reducing memory traffic mainly where Numba’s optimisation lies, or are there other places?” The way I would put it is that Numba is almost C, with some OpenMP-like directives on top. So you optimise a Numba function by thinking like a C programmer, which is what is happening here: you have to think about your memory access pattern, and interchanging the two loop orders changes the performance.

Not because Numba is imitating C, though. Both of them lower to native code for the CPU you are sitting on — Numba through LLVM IR — and at that level it is the memory hierarchy that decides, so any language that gets close to the metal wants the same treatment. What Numba adds on top is that you are not obliged to write everything that way: the array notation is still there, and you drop into loops only where it pays — which also means that, where it does pay, the work is yours.

The leak I described at the start, in its plainest form: f_numpy → f_numba_loops.

The short version first. Numba has something called object mode, and what it means is that if you put @jit on any Python function it will work — but it may work in the way where the compiler says I cannot optimise this, I am giving it back to the Python interpreter.

That is the highlight, and it is the general shape of all three of these tools. You build a certain abstraction, which is a small subset of the Python language, and you say: this I understand, and it can be compiled to something very fast; that I do not understand, and it is not part of my abstraction, so I hand it back. The mechanism is a guard: fall off the guard and the more general thing — Python, in this case — handles it.

The longer version, which is what the slide shows. objmode is that boundary drawn by hand rather than inferred: a hole through the type system for one named expression, so the rest of the function can stay in nopython mode. The example is Kepler’s equation solved element by element with scipy.optimize.brentq, which Numba has no way to compile.

The three timings are the point, and the point is that nothing happened. What objmode buys is a boundary drawn where you can see it: forceobj gives up on the whole function at once, whereas this gives up on one named expression and says so in the source.

Numba is fundamentally type-based. Either you spell the types out by hand, as I did above, or you skip it and at call time it works out what the types are and compiles a version of the function specialised to them. JAX is not like that: it is a tracing compiler, which in a certain sense is similar and in mechanism is quite different.

You do not specify a type. When you pass an array into the function, JAX substitutes a stand-in for it: a Tracer, which carries no data at all, only an abstract value — a ShapedArray, which is to say a shape and a dtype and nothing else. Then, because the Tracer has gone in, the function actually gets called and runs, just as it would on any other object. Drop a print in the body and you can watch it happen. The tracing compiler is observing exactly that: the Tracer goes through all the logic in the function, every array operation it meets gets recorded, and the recording is the program.

So it is even more abstract than Numba. It does not see your source code, it does not even see your bytecode, it does not see how you wrote the logic — it traces the effect of the function. From there the representation goes to a jaxpr, then to StableHLO, then, depending on the architecture, in this case CPU, to LLVM IR and assembly. The windows onto it are jax.make_jaxpr and .lower().compile().as_text(), at the same altitude as Numba’s inspect_types and inspect_llvm.

The interesting thing about JAX, which I have mentioned before, is that it specialises not only on type but also on shape. Feed a 1D float64 array of 1024 elements into a function, then one of 4096, and it will compile a new version for the second. That is where a lot of its advantage in accelerating a function comes from.

Not here, though. Compare this against what we had before and the JAX case is actually much slower — the numpy and numba figures are the same benchmark from a few slides ago, re-run, which is why they are close to those numbers without being identical. Change the array size and it comes out differently: at 1024³ I would not read “best” too hard, because it depends on the run and benchmarking is not that repeatable, but all three are in the same ballpark. This function is also simple enough that you cannot optimise it much further than it already is. In the PyCon UK talk I showed plenty of cases where JAX is much faster — tens of percent, sometimes considerably more, precisely because it can see the shape too.

That prompted the right question from the room: so where is the overhead? Two places. First the tracing and compiling, which you pay on the first run. Second, because it is now dispatching on type and shape, there is a dispatch mechanism to go through: you arrive with some input, it has to find which compiled, specialised version of the function to run, and if that version does not exist yet it has to hand the work back to the compiler. So it is doing heavier dispatching than Numba’s thin typed wrapper, and it is built for a different regime — kernels big enough, or repeated often enough, to pay that back.

Back to the Mexican hat, unchanged, this time given to JAX.

This is the compiler I said existed, the one that can see past a named intermediate. Tracing throws the name r2 away before fusion ever runs, so by the time XLA gets the program there is no name left for fusion to stop at, and it picks its boundaries by cost instead. That is the 127 MB round trip, gone — for structural reasons, not because one fuser is better written than the other. A tighter specification buys more aggressive optimisation.

The price is everything Numba’s fastest version had. There is no JAX equivalent of f_numba_loops and no way to write one, so for a kernel this simple you get the good outcome for free — write the maths, and the compiler does what I had to do by hand — but the cost moved rather than vanished. You pay it in retracing on every new shape, and in what you are no longer allowed to write inside the compiled region: no data-dependent control flow, no mutation. Branch on a traced value and you get a ConcretizationTypeError, which is what jax_fail.ipynb is about.

Everything above is a toy, built to make one point each. The one production kernel I have proper data for is the interferometer curvature calculation from PyAutoLens, which I benchmarked both ways for the PyCon UK talk.

I have made the same mistake in the other direction, which is why I keep saying it. In 2021 I deleted my own Numba kernel from TOAST and replaced it with C++ behind pybind11, on the strength of a ~30% speed-up (toast@a38d1d6). Going back to it properly for this talk, the Numba version I benchmarked against was carrying 893 MB of temporaries per call that it did not need — out += weight * array allocates one, because the subexpression has to be evaluated before the addition can happen — and writing the loop instead removes them in nine lines of Python. Of the rest, a good part is alignment: NumPy hands back buffers that are reliably 16-byte and reliably not 64-byte aligned, so every AVX-512 load straddles two cache lines, and on this machine that alone is worth about a fifth of the runtime. The parallel for simd I was pleased with peaks at two to four threads and is slower than serial by 32, in either language, because the kernel has an arithmetic intensity of 0.25 and nothing left to overlap.

The swap was still the right call — but for reasons that have nothing to do with the number I put in the commit message. The whole thing is worked through in numba_vs_pybind11.ipynb.

And here is that same function again. It looked much the same in NumPy, in Numba and in JAX, and at first glance it looks much the same in Julia too. Three ways this time.

Tidy it up a little and you get those odd dots, which say you want the operation broadcast — and broadcast is the equivalent of what is called fusing in JAX, which I did not properly define earlier, so let me do it here. Imagine you want to go through an array, for i in something, and you want to do it twice: in the first loop you add one to each element, in the second you square it. Fusing means combining those two loops so that you do not have to go through the whole of memory again. That matters when the kernel is memory-bound, which it usually is, because otherwise you keep streaming the array in from memory and going back through it. In JAX it is what made r2 disappear. Here you get it with the broadcasting operator, which has the same effect: do not walk this array over and over creating new intermediate allocations every time.

So the naive function is the first one; do the broadcasting to remove most of the unnecessary I/O and it is more than twice as fast. And Julia happens to have a native function that does exactly this job, so we can throw it in as a third: call mul! and it is faster again, which is no surprise: a linear algebra library is specialised for exactly this, and it never needed a temporary at all. It also shows Numba’s BLAS boundary from the other side: A * B and mul! both become an opaque library call that the Julia compiler, lowering included, cannot see into.

And it is the same story here as everywhere else: you have layers of abstraction in Julia too, with @code_lowered, @code_typed, @code_llvm and @code_native as the windows. The one thing worth noticing is that @code_lowered already shows the fusion decided, before any type is known. Nothing in the compiler chose it. You did, with the dots.

This next one introduces something fairly unique to Julia. It has a type system, which plenty of languages have, and it has multiple dispatch, and combining the two is what makes it powerful. The illustration is Unitful.jl, which, as the name suggests, puts units on your numbers.

Take a function calculating kinetic energy. Call it with plain numbers and you get a number. Attach units and you get a number with units — and note in passing that the u"kg" syntax is itself a bit of metaprogramming, a library inventing new syntax, which is the subject of the next section.

Attaching units to numbers is not new; you can do it in Python. What is interesting here is that a Unitful Quantity is a new type, and for most practical purposes, as far as the compiler is concerned, it is a Float64 carrying a compile-time tag, so it has no runtime cost — and I should nail down that piece of vocabulary, because it trips people up. In a just-in-time compiled language, compile time is the first time the function gets run; every subsequent run of that function is runtime. So at runtime it costs nothing: have a function you call a million times and it is free.

One way to see the effect is to look at the IRs. For a function as simple as this one, the lowered IR of the two calls is identical. The next one down, the typed IR, is near-identical — not exactly the same, because you have a different type. Go further down and the LLVM IR is mostly the same too. That is the proof that there is no performance difference, and the measurement agrees: 1.443 ns against 1.443 ns. The only overhead you might think about is the flip side, that the next time you throw in a new unit you have to compile the function again. Depending on the use case that is free, or it is expensive if you keep changing units all the time.

The other thing I want to mention is that Julia’s array types are just as rich, and I liked that before I had even started using the language: you can have dense matrices, sparse matrices, symmetric matrices, tridiagonal matrices, and so on. All of those are types, so you can expect a matrix to be symmetric, and because of multiple dispatch you can handle each of them differently. That is where composability starts to feel magical: two libraries that have never heard of each other compose, because the methods for those types are already defined. Define a new number type — a quaternion, say — and a differential equation library that only ever knew about real and complex numbers may well just work.

The reason I bring it up is that on the project we ran into a problem here, and it surprised me. I had a sparse matrix, and I did something called ustrip, which just removes the units and gives you back the bare numbers — about as ordinary an operation as there is. Apply it to a sparse matrix and you get a sparse matrix. Apply it to a symmetric sparse matrix and what comes back is dense.

Fortunately I caught that before it went anywhere near production, because it would have been disastrous: these sparse matrices get really, really big, and if I calculated correctly there is no single node that could hold the dense version in memory. It is not a slowdown, it is a program that does not run. And the thing to notice is that it changed the exponent and nothing in the program said so — no error, no warning, no annotation anywhere that changed. The same class of thing bit me trying to feed Unitful quantities to a differential equation solver, which also would not work.

The reason, without going too far into how the language works, is that you still have to implement the methods. Think of it like Python classes: subclass something and you have a subtype, and sometimes the methods keep working automatically, but sometimes what you did to the subclass means that method has to change too. It is similar here. You may need to define a new method to handle that particular case, and it is an explosion of combinatorics — there is some unforeseen combination of things you did for which no method exists to help you. The composability is real and so is the explosion, and they are the same mechanism. The fix is always local, one missing method; the hard part is noticing.

Metaprogramming is a key selling point of Julia.

We did not put it this way when we introduced Numba and JAX, but once you have seen those examples you can view Python, in both cases, as the metaprogramming language — and the Numba or JAX function you are compiling as the mini language. That is not typical. Metaprogramming in C is macros: a very small language sitting on top of a big one that does all the actual work. In Numba and JAX it is the reverse. You have a very small language doing a subset of what Python can do, and a big metaprogramming language, Python itself, above it. Anything you cannot do in Numba or JAX, you metaprogram in Python.

In Julia that boundary disappears, and in that sense I think it is the best JIT for scientific computing. You have a full and quite expressive language — bigger than C — which is also its own metaprogramming language. All code is data.

The example evaluates a polynomial, and the shape of the problem is worth thinking about. For a polynomial of degree you have coefficients. You can throw in any x, and x is what keeps changing between calls; the coefficients are usually constants, and the degree may differ between call sites — I might want a degree-100 polynomial next. Write the function the ordinary way and the number of coefficients is unknown until runtime, which has a performance implication. Can you optimise that away? Yes, and Julia gives you two routes to it, at two different points in the pipeline.

The macro is the one I talked through. You define it so that when it is called it expands into an expression with the coefficients already in it as known constants. That is what goes to the compiler, so the compiler gets to assume the coefficients are constant, and the unrolling of the loop happens for free — look at the expansion and there is no loop left, just one big expression. There is no inner loop left over an unknown number of coefficients.

The other route reaches the same place from the type system rather than from the syntax. A Tuple{} and a non-empty Tuple are different types, so writing horner_step as a pair of methods makes the recursion terminate by dispatch, and type inference unrolls the chain at compile time. Two entirely different mechanisms, the same machine code, the same ~5.3×. Memory is identical across all three versions — one similar(x), nothing incidental — so the whole of the difference is in the arithmetic each one generates.

None of which is exotic: Base Julia ships exactly this as evalpoly/@evalpoly, one function covering both mechanisms depending on whether the coefficients arrive as a tuple or as macro arguments. For contrast, the same idea in Python takes about thirty lines of descriptor-and-proxy machinery and still leaks through the object model in several places.

This last one has no example on it; it is the design of Julia rather than a demonstration of it. Open generic functions: the same function can behave completely differently depending on what you give it, because when you implement a new method with different argument types, multiple dispatch sends the call there. Change the type of an argument to mul and you dispatch to whatever method someone wrote for symmetric matrices. Add aggressive specialisation, and accessible compiler infrastructure — you can program the compiler, because Julia is its own metaprogramming language — and what you get is the power to build more abstraction over what the language can do. That is not something we saw in the case of Numba or JAX.

The one-sentence claim at the top is my own summary, not a quotation from any of the links under it. Why We Created Julia is the original post that launched the language, and it is still a good read; I liked it even before I used Julia. The message is essentially: we are greedy, this is what we want, and that is why we made it — a great many things at once, in one language. Reactant.jl is the one to notice next to it. JAX is powered by XLA behind the scenes, and Reactant compiles Julia functions to target XLA, so if you go that route some of the power of JAX and XLA can benefit your function too — you write the maths and let the compiler do the rest. The GPU backends are all built on GPUCompiler.jl, and KernelAbstractions.jl and JACC.jl are the write-once-run-anywhere layer above them. All of it exists because the compiler is programmable.

That resolves something left hanging back in the Numba section. When an abstraction leaks — when you have to write the host language differently to convey an intent the compiler could not otherwise act on — everywhere else here you get exactly one move: narrow by hand, in the source, once per call site. f_numba_loops is that move, and it is a compiler pass with no name, one you re-derive at every site and nobody else can reuse. @horner is the same pass with a name, written once; and evalpoly is what it looks like once the pass has graduated into the standard library.

The price is the previous section seen from the other side. The seam between the language and the compiler is gone, so you cannot tell by reading which of the two you are looking at: @horner(x, ...) looks like a function call and is a compiler pass, while ustrip.(u"MHz", Symmetric(Su)) looks like one operation and is a lowering decision that changed the exponent. You can build the abstraction properly because the layers are not sealed, and it fails silently for exactly the same reason.

If someone asks which to use, the answer is on the “you write” row. The optimisation strategy is baked into the language design; you do not get to choose it independently of the programming model.

One caveat on the targets row, because the table has to be terse. Julia out of the box compiles for the CPU you are on and nothing else; the four GPU vendors come from the packages in the previous section — CUDA.jl, AMDGPU.jl, oneAPI.jl, Metal.jl, all built on GPUCompiler.jl, with KernelAbstractions.jl over the top if you want one kernel to run on all of them. It is still one source, which is the contrast with numba-cuda, but it is not free of choices.

The line under the table is the general form of that row, and a language decides how much those three sets overlap. Every section above is an instance. In Numba everything compiles — that is what objmode and forceobj mean — but only a subset of it is fast, and that subset is unidiomatic Python: f_numba_loops is C spelled in Python. In JAX rather less compiles at all, but what compiles is fast and it is also the idiomatic thing to have written; you wrote the maths. In Julia nearly everything compiles and most of it is fast, but the sets come apart silently — ustrip on a Symmetric sparse matrix was idiomatic, compiled fine, and changed the exponent.

That is the answer to the question I opened with. You narrow by hand exactly as far as those three sets fail to overlap.

Every section above walked the same staircase, and the plurality of intermediate representations is the reason each of those sections had a set of introspection tools to walk down — and the reason they kept looking familiar from one tool to the next. inspect_types, make_jaxpr and @code_typed are the same instrument at the same altitude in three different systems.

f_numpy → f_numba_loops was semantic narrowing — same language, same result, fewer permitted behaviours — and I did it by hand, because Numba’s specification was not tight enough to do it for me. The dots in mul_fused! are the same thing. JAX does not need either, because it was handed a narrower language to begin with. So the useful question about any of these tools is never “is it fast?” but “which language am I actually writing?” — and each of them is an answer to who does the narrowing: Numba says you, JAX says the compiler, Julia says you and then hands you the compiler so you can automate yourself.

Patch it at the source, or add a rung: that is the general form of the Julia observation above. Only the second response accumulates — a rewrite you repeat leaves the tower exactly as tall as it was.

The last one is the one that matters most where I work. Ship an ahead-of-time binary to a heterogeneous cluster and it has to target the lowest common denominator of every node it might land on; a JIT compiles on the node it is running on, so it gets the right architecture for free. Anyone who has fought module load and a wall of architecture flags to get one binary onto a cluster knows what that is worth.

The last bullet is the one I have actually been bitten by. On an HPC system the compile lag is normally easy to justify, though not simply because the jobs are long: it is that what you are compiling is a long-running numerical kernel, so what you gain in runtime usually outweighs what you spend compiling, by a lot. And it is also where the operational details get in the way. Hundreds of ranks starting at once and all compiling the same kernel against a shared filesystem is not a theoretical problem, which is why the Numba call I ship in TOAST sets cache=False: every one of them writing a compilation cache is worse than not caching at all.

In a compiled language you compile a binary, and that binary is a thing you can archive. It is one of the reasons people like containers — part of what you are doing is archiving all those binaries, and you can inspect them, run static analysis on them, hand them to somebody else. If everything only happens at runtime, just in time, then you do not get those artefacts. And you do not exactly know what it was doing in this instance when you ran it, as opposed to what happened on my laptop.

It is worse than it first sounds, because a compiler will happily optimise your code by specialising on your hardware. Sitting on different hardware can itself change what ends up in the compiled code: a different instruction set, different vectorisation, a different reduction order, different last bits. That is the same problem as the one in my reproducibility talk from last November, seen from the compiler side rather than the environment side.

The previous section is that problem in our own domain, science and reproducibility. This is the same problem from another angle, which is security. If you cannot look at the generated code beforehand, run your analysis over it and sanitise it, then you have a new question: how do you trust machine code that was compiled a moment ago?

That is a real problem on something like iOS or iPadOS, where you have W^X — W for write, X for execute — meaning that a particular region of memory can be either writable or executable, but not both at the same time. Just-in-time compilation does exactly the forbidden thing: read some source or data, generate machine code from it, write that into memory and then execute the memory you just wrote. So on those devices you can almost never run one, with a single exception — the browser’s JavaScript engine, which is itself a JIT compiler and gets the dynamic-codesigning entitlement to say so. WebKit pays for the privilege with a good deal of machinery, the “bulletproof JIT” among it. The point is that the opacity has security consequences, because you cannot see the binary or sanitise it ahead of time.

It is also why there is no Numba on my iPhone, and would not be even if somebody ported it: generating machine code and then executing it is exactly what the platform will not let a third-party app do.

Everything that follows is a variation on the last bullet.

The metaprogramming half is the stronger of the two claims. I do not have to think about the different layers of abstraction between what I want and machine code; I can just tell it my intention. That is my metaprogramming language.

This is the part that most needs saying, because the analogy does not survive contact with it. Everything I like about the comparison is on the previous slide, and all of it rests on a referent that is not there.

Of everything in the digression this is the part I find most useful, because it answers the “but it just hallucinates” objection without having to deny it. Climbing back up the tower means inventing information that lowering destroyed, and invention is the one thing a translator cannot do.

Now the example, and the reason I find it interesting. Bun is a JavaScript runtime — the machinery around the JavaScript engine rather than the engine itself, which is JavaScriptCore — and that machinery was written in Zig. In May 2026 it was rewritten in Rust, using a lot of agents. The diff is humongous — far beyond what a human reviewer can do, and nowhere near fitting in the context window of a single model.

And it got done. In this case the LLM becomes a transpiler, and what made that possible is the test suite. Bun runs JavaScript, so its tests are written in TypeScript — which means they are independent of the implementation language. Change the language underneath, from Zig to Rust, and it does not matter; you have exactly the same test suite. That becomes the verifiable bit. It is almost a spec, in the compiler sense: something the model can transpile against and check itself off. Not for every valid construct of the language, but for what it actually built.

The second half of the slide is the caveat I would put on my own enthusiasm. Nineteen known semantic regressions still got through those 1,386,826 assertions. For half a million ported lines in eleven days that is a good result rather than an indictment, and the interesting thing is where they landed: in the places the tests were silent about. A test suite is a sample of the specification, not the specification.

This is the security section again, taken to its limit. The input space is infinite, so no input is wrong and there is nothing to reject. The output space is infinite, so every output is possible. And under agentic use, whatever the agent decides to do in between is also an infinite unknown space. So you never get to sanitise anything, at any point along the way.

That is the same move as the Julia section, one altitude up. There, the fix for a leaky abstraction was to stop patching at the source and add a rung to the tower. Here the missing rungs are the specification and the guard, and the work is to put them back. Which is why I do not think this is a new discipline: nothing on that list of tests, types, contracts and enforced abstractions was invented for agents, and most of it is already sitting in the repositories we maintain.

If only one line of this survives, let it be the third: the question is not “is it fast?” but which language you end up actually writing, and how much of the narrowing the tool leaves to you.

The one exception to the memory-traffic line is Julia’s unrolling. Memory was identical across all three versions there — one similar(x) and nothing else — so the whole of that 5.3× is in the arithmetic each version generates.

This is the core technical distinction for the session. Willison (2025a) defines an agent as “something that runs tools in a loop to achieve a goal.” The shift from autocomplete to agents is what changed DHH’s mind — he didn’t care for autocomplete (“When I code, I want to finish my own thoughts and sentences” (Hansson 2026)) but agents feel “like working on a team.”

Karpathy’s Software 3.0 framing helps (Karpathy 2026): Software 1.0 = explicit code, 2.0 = learned weights (ML), 3.0 = programming through prompts, context, and tools. The context window is the new program; the LLM is the interpreter.

Willison proposed “vibe engineering” in Oct 2025 and later acknowledged that “agentic engineering” was winning as the preferred term (Feb 2026 update on his post) (Willison 2025c). Karpathy formalized the distinction at Sequoia Ascent (Karpathy 2026). For this session, treat the two as roughly synonymous — both describe the disciplined end of the spectrum.

Karpathy’s floor/ceiling framing (Karpathy 2026): “Vibe coding raises the floor. Agentic engineering is about extrapolating the ceiling. People used to talk about the 10x engineer. I think this is magnified a lot more.”

Some attendees may have only watched Karpathy’s 30-minute Sequoia video (Karpathy 2026), which covers most of this vocabulary. The slides build on that foundation while adding the Willison and DHH perspectives they may not have encountered.

Karpathy pins the step change specifically to December 2025 (Karpathy 2026). DHH independently converges on the same timeline (Orosz 2026). The “macro action” framing is Karpathy’s — the programmer becomes an orchestrator, not a typist.

DHH also names open-weight models in this club: MiniMax M2.1, GLM-4.7 (Hansson 2026).

Note for RSEs: this is very recent. If your last serious encounter with AI coding tools was mid-2025, the landscape has fundamentally shifted. December 2025 was a step change, not a gradual improvement.

This reversal is the narrative centrepiece of the session — the most prominent sceptic went agent-first in ~6 months (Orosz 2026). DHH still cares about beautiful code, craft, and taste. But he barely types code by hand now.

His current setup: tmux with two models running (fast Gemini 2.5 + powerful Opus), neovim in the centre for reviewing diffs via Lazygit (Orosz 2026). He’s not doing vibe coding — he reviews everything, insists on quality, runs tests.

The “mech suit” vs “project manager” distinction matters (Orosz 2026). DHH was put off by the idea of being a project manager of agents. What changed was that agents became capable enough that he feels in control, hyper-accelerated — like wearing a mech suit, not supervising interns.

Could ask the room: has anyone here had a similar shift in the past year?

This is the core of Willison’s “Vibe engineering” post (Willison 2025c). The listed practices are not new — they’re what we already teach as good software engineering. The twist is that agents make them non-optional.

Without tests: the agent might claim something works without having actually tested it, and any new change could silently break an unrelated feature. Without documentation: it hallucinates APIs. Without version control: agent-introduced regressions go unnoticed.

The “amplification” claim is the most important single sentence in the reading set. We’ll interrogate it later. For now, note the implication for RSEs: if your group already does these things, you’re well positioned. If not — if your research code lacks tests, documentation, CI — agents will expose those gaps rather than paper over them.

This slide reflects personal experience rather than the assigned readings. The readings (Willison, Karpathy, DHH) consistently name testing, documentation, and version control as the preconditions for effective agentic work. Reproducibility subsumes and reinforces all three.

Point 1 connects to the functional programming literature. Immutability and the separation of pure and impure code reduce the surface area an agent needs to model. An agent working on a side-effect-free function is operating in a well-defined local world; an agent working on code that mutates shared state must track far more context.

Point 2 is particularly relevant for RSEs: research environments are often inconsistent across machines, making it hard to delegate tasks without extensive hand-holding. A reproducible environment (pixi, conda-lock, Docker, Nix) removes that friction entirely. Connecting to CI: good CI is a reproducible bootstrap loop, already automated and run on untrusted machines.

Point 3 reinforces Karpathy’s verifiability claim (Karpathy 2026). Reproducibility is not just a software engineering virtue; it is a precondition for the kind of feedback loop that makes agents effective. If the same code can produce different results on different runs, the agent cannot distinguish “my change broke this” from “the environment is flaky.”

Point 4 connects back to (Willison 2025b): YOLO mode becomes much less dangerous when the agent’s access is scoped to a reproducible, self-contained project image. The agent can run destructively without touching anything real.

Points 5 and 6 are practical consequences rather than philosophical arguments. Worth mentioning briefly to reinforce the theme: reproducibility is not just about scientific rigour, it is an operational property that makes the agent’s job easier and safer at every step.

Willison (Willison 2025a): “Any time you find yourself thinking ‘ugh, I’m going to have to try a lot of variations here’ is a strong signal that an agentic loop might be worth trying.”

Karpathy’s verifiability framework complements Willison here (Karpathy 2026). “Traditional software automates what you can specify. LLMs automate what you can verify.” Coding is a sweet spot because feedback is immediate and unambiguous: tests pass or fail, code compiles or crashes. Karpathy also warns about “jagged intelligence” — models spike in capability where they are verifiable and heavily trained, and fail surprisingly in areas outside that intersection.

DHH’s P1 optimisation is a good concrete example of a well-designed agentic loop (Orosz 2026). A senior engineer at 37signals asked: what about the fastest 1% of requests — can we make the floor faster? Took P1 from 4ms to under 0.5ms. 12 PRs, ~2,500 lines changed, done in a few days as a side project. This is work that wouldn’t have been contemplated before.

YOLO mode (auto-approving all agent commands) is dangerous but dramatically more effective. Willison’s mitigations (Willison 2025a): sandbox (Docker, Apple containers), run on someone else’s machine (Codespaces), or accept the risk. “Most people choose option 3.” For RSEs handling sensitive research data, the exfiltration risk deserves special attention.

The second point draws on well-replicated findings that using LLMs to answer questions — rather than struggling through them yourself — reduces retention and understanding. This is the same phenomenon Karpathy flags at a higher level: “You can outsource your thinking, but you can’t outsource your understanding” (Karpathy 2026).

The first point is a strong argument for deliberate practice with agents on familiar territory. Most RSEs sceptical of agents have never tried them on problems they already know well. Using a known problem removes the confound of “is this output right?” and lets you focus on “how do I prompt, review, and steer?” That is where the real learning happens.

This heuristic also reframes the on-ramp: don’t start with your hardest unsolved problem. Start with the last bug you fixed, the last function you wrote, the last refactor you did. This is how you build taste for directing agents.

The amplification claim is central to all the readings and worth pushing on. It’s also a claim that conveniently flatters the audience most likely to be reading these articles — and sitting in this room.

DHH reports the same at 37signals: seniors benefit more (Orosz 2026). The survey adds texture: Builders (quality-focused engineers) benefit from AI for large refactors and “quality of life” tasks that wouldn’t have been worth the time before. But they also deal with more AI-generated slop from colleagues, spend more time debugging agent output, and some report identity loss and grief (Orosz and Nilsson 2026).

The Orosz/Nilsson survey identifies three archetypes: Builders (quality-focused, benefit but also suffer), Shippers (outcome-focused, most enthusiastic, risk: faster tech debt), Coasters (can uplevel faster, risk: generating slop) (Orosz and Nilsson 2026). AI amplifies the tendencies that existed before.

The junior pipeline question is economic territory we’re not opening today, but name it as a known gap.

These questions are deliberately provocative and should fuel the discussion. The guitar analogy resonates because many RSEs chose this career partly for the craft of coding itself. DHH abandoned his own analogy within months — he now says reviewing code is where the craft lives, and insists that beauty still matters (“when something is beautiful, it’s likely to be correct”) (Orosz 2026).

Karpathy’s distinction between thinking and understanding is the philosophical anchor (Karpathy 2026). Agents can do the thinking (generate code, explore approaches), but the human needs understanding to direct them — know what’s worth building, what result is suspicious, what tradeoff is acceptable. His MenuGen payment bug illustrates: the agent matched Stripe emails to Google accounts to assign credits, but those can be different emails. Plausible code, bad system design. A human needs enough product and engineering judgement to catch that.

For RSEs specifically: if we value understanding the science behind the software, does agentic coding help or hinder that understanding? The authorship question also has practical stakes for academic credit and reproducibility.

Ira Glass’s “taste-talent gap” is about the frustrating early years of creative work: your taste (what you know is good) outpaces your current ability to produce it. Coding agents invert the usual bottleneck — taste becomes the binding constraint, not raw coding ability.

This reframes Willison’s amplification claim (Willison 2025c): what gets amplified is taste and domain knowledge, not raw coding speed. A researcher with deep physical intuition who never mastered Python can now direct agents well. A competent programmer with no domain understanding cannot.

The evaluation ceiling problem is the converse of the amplification story: once a model’s capability in a domain exceeds the reviewer’s level, the human can no longer reliably distinguish excellent outputs from merely plausible ones. For RSEs, this is a near-term practical concern in specialised subfields — and it makes domain expertise the durable investment.

This slide is reference material only. We deliberately excluded the economic dimension from today’s discussion — the readings were selected to focus on the technical, psychological, and philosophical dimensions instead. If attendees raise cost or labour-market concerns, acknowledge the omission rather than improvise.

The METR study (not assigned reading) found that experienced open-source developers believed they were 20% faster with AI tools while being measurably 19% slower. Worth mentioning as a provocation if anyone insists on the productivity question, but do not open it as a full discussion thread.

This is the bridge from the readings to the audience’s working context. Pause here and let the room sit with the question before moving to the canned questions.

Research software often lacks the preconditions that Willison and Karpathy assume: comprehensive test suite, clear success criteria, stable specs (Willison 2025c; Karpathy 2026). In research, “correct” might mean “agrees with physical theory,” “matches experimental data within error bars,” or “reproduces a known benchmark” — none of which are trivially automatable. The tension is sharp: the readings all agree that tests are the foundation of effective agentic coding, but RSEs often work in domains where the test harness is the unsolved problem.

You probably won’t get to all five. Pick one or two that the room gravitates toward. Question 1 is a warm-up that also surfaces how many attendees have hands-on experience. Questions 2–3 probe the core claims. Questions 4–5 are RSE-specific and likely to generate the most productive discussion.

If time is short, prioritise 4 or 5. Question 4 connects to academic integrity and reproducibility — familiar ground for RSEs. Question 5 connects to practical team management — the tension between individual tool preferences and project-level constraints, which is likely to surface real experiences.

Naming what’s out of scope prevents drift into familiar but unproductive territory.

On “does AI actually make us faster”: the METR 2025 study (an RCT with experienced open-source developers) found a striking mismatch — developers believed they were 20% faster while being measurably 19% slower. Mention this if someone insists on the productivity question; do not open it as a full thread.

On “which tool”: tempting for a technical audience but will consume all available time and produce no actionable insight in a one-hour session.

These are not unimportant questions. They’re just not today’s questions.

Included for reference during discussion. Some attendees may have only watched Karpathy’s 30-minute video, which covers most of the conceptual framework (Software 3.0, vibe coding vs agentic engineering, the December inflection, verifiability). The slides are designed to fill in what Karpathy doesn’t cover: Willison’s specific practice framework, DHH’s personal reversal story, and the Orosz/Nilsson survey data.

Defining reproducibility is surprisingly challenging, as the terminology varies significantly across different fields and communities. When researchers discuss reproducibility and related concepts, they may use the same words with entirely different meanings, so it is essential to be aware of whom you are speaking with. In our field, we typically focus on what might be called “methods reproducibility,” which centers on the recipe—how you take the same piece of data and use the same software to reproduce exactly the same kind of result. This contrasts with “results reproducibility,” which is more about replicating the science: given the same data, researchers are free to implement their own software to verify findings.

When we talk about computational reproducibility, there are actually many interconnected components. At its core, we are asking whether the same data can produce the same results. This involves not just the data itself but also the software, the workflow (ensuring that complex procedures are reproducible and produce the intended outputs), and the operational environment. Binary reproducibility—where compiling the same source code twice produces bit-for-bit identical binaries—is technically achievable but extremely difficult and, for many practical purposes, not critically important. Source reproducibility, where the same program with the same input produces exactly the same output, is also challenging because factors like floating-point arithmetic being non-associative mean that adding numbers in a different order yields different results. The focus of this discussion is primarily on environment reproducibility: how can you ensure that from a given source code, you can build the software and construct an environment (potentially including multiple software components, dependencies, and even the operating system itself) to load your data and ultimately obtain your results?

The term “package manager” can refer to several distinct but related concepts. First and most obviously, it refers to the software program itself—the command-line tool you invoke when you type conda, pip, apt, or similar commands. But when you tell someone they can use conda to install SciPy or NumPy, you are not merely describing the capability of the software; you are implicitly referring to the ecosystem in which that package has been prepared and made available. The packaging work that enables this installation is substantial.

Additionally, package managers involve an index (sometimes called a channel or registry). For users of systems like Ubuntu, adding a repository URL before installing certain packages is a familiar operation—you are telling the package manager where to look for the software. Finally, a package manager can sometimes refer to an entire distribution. For example, when people install Anaconda, they are getting a distribution: a bundle of the conda command-line tool, a default channel, and a curated selection of packages. Similarly, Miniforge is a distribution containing conda, mamba, the conda-forge channel, and a minimal set of packages.

Package managers can be categorized in several ways. Some are specialized for operating systems—apt for Debian-based systems, for instance—and you would typically not want to use these for building your entire scientific software dependency stack. However, you would need them to install low-level, system-specific components. Others are language-specific, such as pip for Python. There is also a distinction between source-based and binary-based package managers. In the Python ecosystem, historically (and sometimes still today), packages could be distributed as source distributions, requiring the build process to occur on your local machine. Some package managers exclusively fetch source and always build locally, which can be useful—particularly in the HPC context where you want software optimized for a specific architecture. Binary-based package managers, by contrast, distribute pre-compiled packages. Platform support is another important consideration: can the package manager handle different computer architectures or operating systems? Cross-platform capability is valuable for deployment scenarios where you package your software once but deploy it across multiple environments.

With this background established, we can now articulate the specific problem statements this discussion aims to address. First, how do we reproduce an environment containing a bundle of research software across different systems and platforms? “System” here refers to operating systems like macOS or Linux, while “platform” refers to architectures like x86-64 or ARM. Second, how do we customize the build process—for example, using vendor-provided compilers or interconnect libraries like MPI—while still maintaining reproducibility? This is particularly important in the HPC context, where we need to deploy software and want to extract maximum performance. Third, once we know how to build an environment on a particular system (whether a local computer, a workstation, or a specific HPC cluster), how do we translate that environment to another system?

This example from the pixell package illustrates the problems inherent in PyPI packages. First, note that this is a setup.py file, which is notorious and increasingly deprecated—the package has since moved away from this approach. But this serves as a useful illustration of how problematic things can become. The script checks whether it is running on Windows (which it does not support), and on Darwin (macOS) or Linux, it attempts to infer what compiler is available on the system. At various points in the script’s history, it would list executables in different paths like /usr/bin or /usr/local/bin and try to locate compilers such as Fortran compilers. If unsuccessful, it would search other paths. This approach is inherently unstable and completely non-reproducible. If someone can install this package on their computer, there is no guarantee that it will successfully compile on another computer. The error messages in such cases can be extremely cryptic and difficult to debug.

This quote comes from Guido van Rossum himself at a 2012 meeting. The statement essentially acknowledges that scientific computing needs are so unusual compared to the larger Python community that building a specialized package manager makes sense. This was the genesis of conda. The developers had been experiencing numerous build problems with complex scientific software—packages involving C compilers, Fortran compilers, and intricate linking requirements. The complexity became so overwhelming that this statement became the catalyst for creating conda.

The conda solution addresses many of these challenges. One key aspect is the separation of build dependencies, which isolates your environment so that builds become more reproducible and less dependent on the host machine. This provides a higher degree of reproducibility: if you install something on your computer and it runs, there is a high probability it will work on another computer as well. It is not guaranteed, but the probability is significantly higher than with traditional approaches. The concept is sometimes called a “hermetic build”—the ability to bootstrap everything from scratch. Conda is close to achieving this, though not 100% hermetic. This is a useful property for achieving reproducibility.

Conda supports many platforms, though notably not all—OpenBSD and FreeBSD, for example, are not supported. From the ground up, conda was designed to be language-agnostic. While it is primarily associated with Python packages, you can use it to package anything: C/C++, Fortran, R, Rust, bash scripts, and even Julia (via juliaup). Though it should be noted that Julia’s packaging situation through conda is not ideal, this is not conda’s fault.

There is some potential confusion around the terminology in this ecosystem. Mamba started as a new solver borrowed from Red Hat Enterprise Linux to overcome the performance problems of conda’s SAT solver, which was implemented in Python. The original conda solver works correctly but is very slow—as people built larger and larger environments, solving could take hours. Mamba has now fully matured into an almost drop-in replacement for conda, and its solver has been merged back into conda, so the default conda solver is now also fast. Personally, I now only use mamba.

The default channel in conda/mamba points to packages built and maintained by Anaconda. However, anyone can create new channels hosted on anaconda.org. The most prominent community channel is conda-forge, which is many things at once: it is a GitHub organization with many repositories called feedstocks. Each feedstock contains at minimum a meta.yaml file that defines how to build the package. Continuous integration is deployed using infrastructure defined by conda-forge to build the same package across platforms, Python versions, and various variants (such as different MPI backends). This is a large and useful ecosystem, and for many scientific packages, conda-forge does a better job of packaging than the alternatives.

If you manage to package your software for conda-forge, it achieves a high degree of scalability and reproducibility. This in itself is a valuable exercise. Personally, when releasing a package, I always release it on PyPI first—partly to claim the namespace, as this is the authoritative source for how people find Python packages. But the next step is to distribute it via conda in some form. Packaging for conda-forge is harder than for PyPI, but this difficulty is the cost of guarantees. The separation of build, host, and runtime dependencies ensures clean builds across different platforms. The CI system is particularly useful because it operates in isolated environments independent of your local computer, testing your package on various systems you might not otherwise have access to.

This addresses the second problem statement regarding customization. You can choose different BLAS implementations—the build matrix is quite large for some packages. By default on x86-64, you might get MKL, but you can also select OpenBLAS, BLIS, Accelerate (on macOS), or Netlib. Different MPI implementations are also available: OpenMPI, MPICH, or no MPI at all. The reason for the “no MPI” variant is that MPI libraries can be difficult to run in certain situations, so if you only need serial execution, avoiding MPI dependencies simplifies things.

Even microarchitecture optimization is now possible. On x86-64, there are officially four different levels of microarchitecture, each with different instruction sets and vectorization widths. AVX-512, for example, uses 512-bit vectors, so with 64-bit floating-point numbers you can process eight operations concurrently in one cycle. Optimization per microarchitecture is important for squeezing out the maximum FLOPS from your CPU, and conda-forge now provides this level of control.

The last problem statement concerns distributing environments across systems. The most rudimentary approach involves exporting an environment to a YAML file—essentially a lock file—and then creating an environment from that file on another system. In some cases, you can reproduce exactly the same environment, though platform differences may prevent this. For example, if the exported file includes MKL libraries (which are x86-64 specific), attempting to recreate the environment on ARM would fail. But within the same architecture (such as Linux x86-64), this approach works reliably.

Cloning environments is another option that might seem trivial—just copying a directory to another path. However, if you understand prefix paths and how they work, simply moving a directory would break your environment. Conda employs a “relocatable” trick that allows environments to work with different prefixes, which is quite useful.

For simple cases, particularly pure Python projects that are already on PyPI, there is a tool called Grayskull that can generate a conda recipe automatically. It will create a recipe for you that is not 100% correct all the time—sometimes minor modifications are needed—but this is as good as it gets for simple projects. For pure Python packages, it often just works out of the box.

This example from the ducc0 feedstock (a CMB-related scientific package) demonstrates more complex packaging. It shows the concrete implementation of concepts discussed earlier: the separation of build, host, and runtime environments. Writing these recipes can be challenging, especially when starting out—you might encounter cryptic errors and need to experiment. The CI allows you to test your package on conda-forge’s infrastructure. Most packages include simple tests like importing the module successfully, which is not a full test suite but provides basic validation that the build succeeded.

This even more complex example from the TOAST feedstock demonstrates how to pin to compiler components without specifying implementation details—dispatching between clang and LLVM, for instance. This is a template defined by conda-forge, not native to conda itself. You can see conditional logic embedded in what appear to be YAML comments (using Jinja2 templating). The recipe handles different BLAS implementations and MPI variants, showing how the build matrix can become quite sophisticated for packages with many optional dependencies.

This example addresses how to reproduce a set of system software on HPC. Consider tools like fzf, fd-find, curl, bat, ripgrep, and others that you might want available on any system you work on. In the past, HPC environments were often quite minimal, lacking many common utilities. I have experimented extensively with different methods of installing command-line tools in user home directories on Linux—from source compilation to Gentoo’s Portage—but this approach is the best I have found. It leverages the fact that conda is a cross-platform, language-agnostic package manager. It does not have to be for Python software; many of these packages are essentially just downloads that get expanded, but the process is more convenient and more secure than manually downloading and extracting archives.

This real-world example from the SO:UK Data Centre demonstrates how to distribute a complex scientific software environment on a heterogeneous HPC cluster. The approach builds on conda to bootstrap Python packages and build tools, with optimization for different microarchitectures. A CI/CD workflow using GitHub Actions creates distributions with predefined prefixes, which are then deployed via CVMFS. This addresses the challenge of maintaining consistent environments across diverse computing resources.

Let me start with why I use it.

I am a cosmologist. I did my PhD in the UC Berkeley Physics department on the topic of the Cosmic Microwave Background (CMB) radiation — the relic radiation that is everywhere, so that every time you look at the sky you are looking at it, you just don’t know it. During my research I used the NERSC HPC facility in the US, a top 10 supercomputing system, to perform CMB data analysis. HPC stands for High Performance Computing, also known as supercomputing.

Briefly speaking, as we cannot do an experiment cosmologically, we instead observe a huge amount of data and use statistical methods such as Bayesian inference to deduce information about our universe, such as deducing the dark matter composition, or finding evidence of inflation during the Big Bang. My main interest in Python is therefore primarily its application to scientific computing on HPC.

This was a question to the room, so treat it as one here too. By “can you use”, I mean a language that has demonstrated it scales to a state-of-the-art, full-system supercomputer — not merely a language you can invoke on one, the way you would invoke bash.

The room got Fortran and C/C++ immediately. The next answer is Python, and you might call that cheating a little bit, because many of the codes you call from Python to run on an HPC system were written in Fortran or C and are simply exposed as a module. It is still very useful as the glue that holds everything together. Julia belongs on the list as well. MATLAB and Rust do not — not because they are bad languages, but because they have not yet demonstrated that they scale to a full system.

The opening keynote resonates with me a lot. Python has superpowers. I recommend people watching that to see the history and the bigger picture. Here, I am more narrowly focused to tell you why, at this particular moment, Python is the language of choice to write and run scientific applications on supercomputers.

The room offered embarrassingly parallel workloads, and developer efficiency. Both are true — and even when the problem is not embarrassingly parallel you can still drive it from Python. But my own answer is this.

First of all, it is where the community is. NumPy, SciPy, Astropy, you name it. I am not claiming to explain how it got here, or whether it will still be the case in the future; only that at this moment, this is where the community is.

Secondly — and I don’t know whether this is a blasphemy — it is because we are the people who don’t care what language our algorithm is implemented in. CPython itself is written in C. C++ is exposed as a Python module through pybind11. Many SciPy functions are Fortran underneath. Write something in Julia, or in Rust, and you can still expose it with a Python interface.

In other words, Python is the glue. If we look back to the list of languages that has demonstrate they can scale to the state-of-the-art supercomputer at a given era, there is not a lot on the list, and Python is one of them. You might call that cheating, because in this case Python is mostly just calling C modules or libraries written in Fortran. But we don’t care. Python is the glue that successfully composed very complicated scientific workflow together.

So, why JIT? To answer that I need to give two short introductions first.

Back to the languages, focusing on the aspect of how source code eventually runs, it falls into two categories, compilers and interpreters. And within compilers, there are also ahead-of-time (AOT) compilation and just-in-time (JIT) compilation.

Here we will name a few examples. The reference Python implementation, CPython, is an interpreter, which is AOT compiled by C compilers such as gcc and clang. Bash is another commonly used interpreted language, and I would say that in the HPC space, Python is increasingly taking over what bash used to do. AOT compilation is what traditionally excels at HPC: C, C++, Fortran.

Another well known Python implementation pypy is a jit compiler — supposedly any valid Python program can run under it. Julia and JavaScript have JITs too.

The second introduction is how you accelerate a piece of Python code when you need to.

I hope we all agree the CPython interpreter is slow, right? To accelerate applications running in CPython, there are multiple frameworks to do it. The traditional model is C modules: write your performance critical part of the code in C and use the CPython API to expose a Python interface to the users. A prominent example of that is numpy, which defines an interface often used in high performance numerical operations, and is architected in C with a pythontic interface. This is still one of the fastest and pythonic way to write numerical code in Python without introducing further compliation to the users. Other examples on AOT are cython and C++ with pybind11, which I won’t go into details.

Turning our attention to the spotlight of the day: JIT. pypy is a general purpose Python jit compiled runtime for a long time. There is also an experimental JIT compiler in CPython 3.13+, which would jit compiled a subset of hot code: anything valid still runs, but only some of it is actually compiled, and it is not aimed at numerics.

The remaining JIT compilers are the one we are focusing on today. Numba and JAX can both be regarded as Domain Specfic Languages that will jit compiled a subset of the Python language, focused on numerical compuations.

My headline answer to “why JIT” is that it solves the two-language problem: you no longer need two different languages to implement very high performance code.

To briefly give an idea why JIT may be useful, we can look at this function as an example. If you are familiar with Numba enough, as soon as you write down this simple function with Numpy, adding the jit decorator there is a no-brainer. You would be able to immediately predict that it will be faster — three times faster than the pure NumPy version, in this case — and that memory allocation is simpler, and hence lesser memory pressure and lesser chance of getting killed by out-of-memory (OOM) error.

However, when I ported this to C++ with pybind11, 14 files are changed, 230 lines are added. If you look into the commit change, for sure a lot of them are boilerplates. But that is exactly the point. It went from using Numba as a 10s decision to being 30% faster after a ton of work. One single line gave three times; all of that work gave 30%.

This sort of things happens all the time when writing science code. How your program eventually got there is often not obvious in the beginning. It involves repid prototyping and sometimes solving complicated mathematical or algorithmic problems. Having a DSL within Python that can be jitted is a superpower Python given us to move fast and run fast at the same time.

Now it all makes sense, kind of. Let us do an actual example to see how JIT works.

I call this a slow track. It is a live demo without being live. I hope it will gives you a feel of the languages and gain intuitions yourself. Later on, we will have a fast track and you just have to believe the high level summaries I am going to give you.

Let’s stare at this equation and see how you would implement it.

The function is not really very long — it is almost a one-liner. I have just given it a lot of whitespace to make it clearer what is happening, but it is only implementing the math as it is written. And it is reasonably fast.

Ignore the tiny little details there where the order of the dimensions in and are opposite to each other.

Mine looks a little complicated, but the part that really makes the difference is parallel=True. It is a very simple decorator to add, and you can immediately see that it is faster.

Can you guess? Just change the decorator — plus one more thing. Plain NumPy does not work inside a JAX-jitted function, but JAX ships its own module, jax.numpy, which is almost a drop-in replacement: wherever you had np, you use jnp, conventionally imported under that name, and it works.

You would notice the speed up from Numpy to Numba to JAX, in this order, which is often the case. There is a general trend of progress here: you start with something in NumPy, you implement it in Numba and it is faster, and then in JAX — where it is not even a reimplementation, just a change in how it is compiled — and it is faster still. That is the general trend, though I am not saying this is always the case.

Ok, case closed, right?

Hold on, not so fast. Does anyone spot any potential problem with this implementation?

All three implementations are the same in this respect, so here is a hint: go back to the NumPy listing and look at the comment marking (M, M, K). What it means is that an array of that shape is materialized at that point in the expression.

I haven’t told you yet how large the problem is. is an matrix, where is the number of image pixels and can be as large as 70,000. is the number of so-called visibilities, and can be as large as . So if you work out the size of the array being created at that intermediate step, it is on the order of hundreds of PiB. There is simply no computer that will fit that in memory.

The answer itself is fine — an array of 64-bit floats is about 40 GiB. The problem is not the answer, it is the intermediate.

The solution is simple, we do not want to expand the -dimension fully in memory. We want to evaluate that lazily and accumulate it over .

How would you do that in NumPy? You can always write a Python loop over — but would that be efficient? Probably not. There is no efficient solution in numpy however without using Python loops. I think there isn’t one; I can’t prove it. Let me know if you know a way.

In Numba, however, we can have efficient loops, we can even do it in parallel. This function looks very similar to the version before: that one line is computing the same equation for a fixed value of , you loop over , and you do an in-place reduced sum. That is all. It will handles the reduced sum there in parallel correctly for you. It actually maps very nicely to C with SIMD and OpenMP parallel-for there. In fact, OpenMP is one of the threading layer backend in Numba.

A large part of why it is this fast is numba.prange, which parallelizes the loop.

This means we can does the same in JAX, right? Wrong. It is very tempting to just swap the decorator — that is the whole promise of a decorator, that you put it there and it works — but here it won’t.

This is where one of the core design of JAX matters. JAX is essentially designed with pure functional programming paradigm in mind. It enforces immutable data, where the in-place reduced sum there is prohibited. JAX dislikes anything that mutates your state.

It does not mean it does not have semantics for looping. To express exactly the same logic, you refactor it in a certain way: a kernel function evaluating the same equation for one particular value of , and then a way to tell the program to iterate over all and sum them up. Here, I choose scan, one of the idiomatic JAX interfaces for that, and itself a functional programming construct. What it does basically is to focus on each term for a dixed , and then accumulate the sum over different values of . If you believe me, this represents almost exactly the same thing as the Numba version, in the JAX-idiomatic way.

See python-autojax/experiments/reduced_sum at main · ickc/python-autojax for more details on how this can be expressed in JAX.

This section is the fast track, and it is a bit denser than the rest. On the day I was short on time and skipped some of it; the article does not have to.

Now that we have seen some examples on how jit works, we will introduce them once more in details.

Numba is a jit compiler supporting a subset of Python and NumPy operations, powered by LLVM. While it is possible to target the GPU via CUDA, it requires rewriting the function using different APIs and paradigms, not to mention it is CUDA (i.e. NVidia) only.

JAX is a jit compiler, tracing compiler by Google, powered by XLA compiler, originated from Google. XLA stands for Accelerated Linear Algebra, which already tells you where its focus lies. JAX is designed primarily for machine learning workloads but is suitable for scientific computing as well, because the two have a lot in common. It is a tracing compiler removing side-effects of function. I.e. effectively it encourages functional programming paradigm and thinking. It automatically targets multiple hardware architectures including CPU, GPU, TPU, without requiring rewriting. I.e. it solves the “two-language problem”, or more accurately, “three-implementation problem”: prototype/API, CPU, GPU.

This is where JAX shines over other solutions to develop for the GPU. In principle, one single implementation is all you need. You do not need to optimize a function for the CPU, optimize another for the GPU.

One more tip for thinking about either of them: it is better to regard the thing named “Numba” or “JAX” as a language, a compiler, and a library all at once. A domain-specific language, because not everything can be jit-compiled; a compiler that does the compiling; and a set of libraries that come along with it. The comparison below is organized along those three lines.

As a language, Numba is more like a C-like mini language: mostly anything you can express in C can be expressed in Numba. JAX is a smaller language still — more restrictions on control flow, on mutation, and on dynamic shapes; you basically cannot do any of those.

As a library, Numba implements a subset of Python and NumPy with a parallelization model of its own. NumPy in Numba is a drop-in replacement, but only a subset of those operations is implemented, and whenever NumPy is called within a Numba jit-compiled function, it is completely hijacked: it is not really calling NumPy, Numba takes over and does its own thing. JAX, on the other hand, implements a subset of Python, NumPy and SciPy, so it is the larger library, and it lives in its own namespaces, jax.numpy and jax.scipy. The APIs are similar to NumPy and SciPy — most of them are the same — but they have their own documentation, and that documentation is precisely what makes it practical for JAX to deviate in behavior where its own restrictions, such as the functional paradigm we mentioned, require it.

A few characteristics are worth internalizing. Because JAX is a tracing compiler that recompiles per shape change, anything that is genuinely a compile-time constant — a shape, for instance — has to be marked as such with static_argnums, or you recompile on every call.

More broadly, I think of this as compiler driven design. The upside is that performance improvements can arrive for free through compiler improvements, so long as your code is written the idiomatic way; the downside is the wall you hit when you want low-level control. It is a trade you make knowingly.

If you find yourself wrestling against the language, deviating from idioms and best practices, time to choose something else!

These are DSLs. They are small on purpose, and the moment you are fighting the smallness rather than using it, the honest move is to step outside: C++ with pybind11 if you want ahead-of-time compilation, Julia if you want a JIT.

Everything so far has been leading to why I actually did this work. The case study is PyAutoLens: a project with a very complicated science codebase, originally implemented in Numba, and I am part of the effort that brought the whole thing to JAX.

We previously have seen different implementations of and some informal benchmarks.

Here we are seeing the benchmark generated using the actual library I wrote.

The numba and jax in this table corresponds to the low memory implementation we’ve shown previously. We see that given 1 single CPU core, the JAX version is slightly faster than the Numba version.

But more importantly, we see that there is a _compact version that I have not shown before. This is 3 orders of magnitude faster. This is done by taking advantage of the fact that δg there has a lot of repeatitive values. The algorithm that I shown to you is while the compact version is .

This also illustrates that sometimes you can focus on low level optimization and have 30% speed up. Or you could spend your time instead on the mathematics behind it and find a better algorithm that can achieves unlimited amount of speed up.

Now how about the million dollar question, how would JAX perform on the GPU? Since the Numba implementation cannot run on the GPU, let us give it a boost to have some resembance of fair fight, 128 CPU cores, vs. JAX on the NVidia A100.

(Notice the slight change of problem size here.)

What we see is the Numba version scales almost perfectly with 128 CPU cores. The JAX implementation is however another order of magnitude faster.

What happens for the compact case though? I think it is because they are so efficient so that the problem size is not big enough to “warm up” the GPU. But note that it is not a fair fight anyway, as we can talking about comparing benchmark from different hardware.

Just to give you one more benchmark to look at, here is another complicated example I have done in the project. You are already familar with , where as here can be a sparse matrix.

While I am not going into details here, but there is a bunch of combinatorics here. Should we calculate directly, or go through the compact version? Should we use the sparse structure of matrix , or expand that in memory first?

My claim is that this depends on the size of input data, and which is the fastest to choose from depends on your data which can be informed by benchmark like this.

Looking at how JAX runs on the GPU, surprisingly the simplest implementation is fastest. Again I suspect there is not enough data for those faster algorithm to flex its muscle.

If you have to pick between the two based on my experience, this is where each one lands. It is not the only thing you can do, but it captures the shape of the decision.

There is one situation where Numba is genuinely better than JAX: if you only target the CPU, you know you have a lot of CPU cores available, and your problem is not linear-algebra heavy. In that corner Numba wins. Most of the time, though, JAX is quite good.

Now, let’s see what JIT has enabled our science.

I am not going to give these graphs a justice. But, I am going to have a crash course here and jump to conclusion. The , etc. you have seen before is part of a calculation known as the likelihood function. The is a big part of turning the observed data into a likelihood function. Now you can make predictions from models given some parameters. and the likelihood function is going to tell you how likely those parameters are. Using different kinds of samplers, which all involve running the likelihood functions repeatedly, often with of iterations. And the parameters that corresponds to the maximum likelihood is simply called the maximal likelihood estmator (MLE).

This is basically how we do inverse problem statistically. And obviously it is some kind of maximization problem.

There is one thing about JAX I have not told you yet, and it belongs here. If you are looking for a peak, gradient information is obviously useful: the gradient already tells you roughly where the peak is, pointing you in the right direction. And now that things are jit-compiled, consider what compilers actually do — they transform code. You have code written one way and it is transformed into another. So why not leverage that to transform your function into its gradient as well? That is exactly what automatic differentiation, jax.grad, does, and in a science case like this it is a real advantage.

Here, we will see a concrete example of how PyAutoLens does it. It is a bit dense, but focus on the image on the right: this is your observed data, and it is an image of a strong gravitational lensing. The question to ask is, out of all those complicated parameters, how would you come up with a model that predicts exactly that pattern?

The animation shows how iterations of different parameters converge on that pattern.

At the start, the image in the top right hand corner does not really resemble the observation. But as it iterates more and more, it gradually becomes it. Eventually, comparing the image in the top right hand corner against the data, we can see the modeled image pretty much ressemble it.

Here is how we can infer on the mass of the dark matter we cannot see.

One more thing you can notice here is, while probably I can implement the likelihood function in any language I want, eventually I need to expose it in a Python interface because I need to pass it to these libraries: Nested sampling (Dynesty), MCMC (emcee), particle swarm optimization (PySwarms).

This is what I mean when I say the superpowers of Python is it is where the community is, and why which language we implement it does not matter. Python is still our favorite glue. With JIT in Python like Numba and JAX, we can stick to the glue a bit longer.

So in this particular case — PyAutoLens — who won the battle of JIT? Can you guess?

JAX, of course. Our project was funded to move things from Numba to JAX; if JAX had lost, we would basically be a complete failure. The interesting question is by how much. The room guessed five times.

It is 50x, according to preliminary results. On the slide that number is set fifty times larger than the text around it, which is roughly the point — I don’t feel good about the typography, but that is how much.

Lastly, I just want to point you to the team of people working on this and some links to share. Thank you.

The project’s motivation is to probe the nature of dark matter with the James Webb Space Telescope (JWST). The science case is to investigate the low-mass region — below roughly — of dark matter substructures in our universe. This is the region where the well-established ΛCDM model and alternative dark matter models make statistically different predictions.

Substructures in that mass range are generally invisible, so we cannot observe them directly: we have to probe them via strong gravitational lensing. This was first demonstrated two years ago in James W. Nightingale et al.¹ using Hubble Space Telescope (HST) data, so the natural next step is to take advantage of the 4+ wavebands and the high quality observations from JWST.

To do this, we use an existing package called PyAutoLens. It has two halves:

In cosmology, we cannot perform experiments; we only have one observable universe. All we can do is make observations and then use simulations to infer the parameters of physical models through statistical methods. That is the loop in the diagram above: physics gives us parametric models, models give us predictions, the cosmos gives us observables that our pipelines compress into data, and a sampler ties the two branches together through the likelihood function.

PyAutoLens uses a sampler to iterate and find the best solution. The goal is to find the mass profile that produces the observed image.

However, there are computational challenges.

That extra cost buys us something, though: because JWST provides data across multiple frequency bands, we get colour information that helps guide the model to better solutions.

The question is: why port PyAutoLens to JAX for calculating the log-likelihood function?

The documentation point matters more than it looks. Numba’s documentation on its internal workings can be minimal, whereas jax.numpy and jax.scipy have their own comprehensive documentation — which is what makes it practical to deviate from the reference implementation when you need to.

Let’s look at some benchmarks. The function under study here is , a weight matrix used in the interferometer likelihood:

Following the methodology, the first JAX attempt is a direct, fully vectorized transcription of the math:

Expanding the cosine of a difference into a product of cosines and sines turns the whole thing into two matrix multiplications:

In other words, the naive vectorized implementation from the 1st try is computationally infeasible — not because the answer is big, but because the intermediate is. The final result is manageable, so the problem is tractable as long as you avoid ever expanding the largest dimension in memory.

In Numba, the answer is to loop over and use prange for a parallel reduction, much like OpenMP:

JAX has no prange, and its arrays are immutable, so an in-place reduction is not available. The idiomatic solution is jax.lax.scan, building up the sum iteratively:

On a single CPU core, the JAX implementation is slightly faster than the Numba version — about 13%. But look at what changing the underlying algorithm buys you: the compact formulation is a factor of four faster than the best implementation of the original algorithm, and over a thousand times faster than the naive ports of it, all without leaving Numba. Algorithmic improvements can be far more impactful than framework changes.

For the original algorithm, JAX on the GPU is significantly faster than Numba on 128 cores — as you’d expect. But for the alternative, compact algorithm, JAX was slower, likely because the problem size was not large enough to saturate the GPU.

On a single CPU core, the fastest algorithm uses sparse matrix operations in Numba.

On the GPU, however, the dense JAX implementation is the fastest by a large margin — and the sparse implementations that won on a single core are now an order of magnitude behind.

Improving the algorithm within Numba can often provide significant speedups without needing to switch frameworks at all.

It is a powerful language with a pure functional paradigm that makes it easy to achieve great speedups and deploy on GPUs. But it is not a silver bullet for all HPC use cases, particularly those relying heavily on multi-core CPU parallelism. Thank you.