{"slug": "torchwright-compile-computation-graphs-into-vanilla-transformer-weights", "title": "Torchwright: Compile computation graphs into vanilla transformer weights", "summary": "Rob Halsmith created Torchwright, an open-source compiler that produces exact transformer weights for any computation graph defined in Python, enabling execution of algorithms like binary increment on a standard Phi-3 model with zero training. The tool compiles primitive operations into vanilla transformer weights, demonstrating that transformers can express algorithms without learning, and is available on GitHub.", "body_md": "For a while now I’ve been chasing a question: what sort of algorithms is a transformer even capable of expressing? I got here by wondering about arithmetic — the first generation of LLMs were really bad at it, and I wanted to understand whether that is a limitation of training or a limitation of the architecture itself. After all, transformers are Turing complete — at least in proofs that idealize the arithmetic and let the model generate as many tokens as it needs along the way. Even so, naively one might expect elementary school arithmetic to be within reach.\n\nI decided to try a specific opinionated approach: instead of asking what a transformer could learn, I would calculate the exact weights necessary for a transformer to implement an algorithm explicitly. If the weights exist, the architecture can express the algorithm, and the question of training never enters into it.\n\nWhat I found when I pointed the finished tool at arithmetic is the next post. This post is about the tool I created to get there — and the tool’s output, concretely, looks like this:\n\n``` python\nfrom transformers import pipeline\n\ngenerate = pipeline(\"text-generation\", model=\"binary_increment_hf_bundle\")\nprint(generate(\"1011\\n\", return_full_text=False)[0][\"generated_text\"])\n# 1100\n```\n\nThat is an ordinary [Phi-3](https://arxiv.org/abs/2404.14219) checkpoint incrementing a binary number, loaded by\nvanilla huggingface with no custom code. Every weight in it was computed. Zero\ntraining. The rest of this post is how.\n\n## The plan\n\nLiterally constructing the weight matrices by hand would not be feasible, so I came up with a plan. If I could define primitive operations that can be easily expressed in transformer weights, then maybe I could compose these primitive operations in increasingly complicated ways. I would then build a compiler, which would create an empty transformer and decide how to allocate the residual stream to all of the intermediate values, and thus where in the various weight matrices each operation belongs. If I could get that foundation in place, I could compose the primitives into more and more sophisticated operations: a library that would make it straightforward to express a real algorithm in a transformer.\n\nThe idea of hand-built transformer weights was not new.\n[RASP](https://arxiv.org/abs/2106.06981) defines a language whose primitives\nmap onto transformer sublayers, and [Tracr](https://arxiv.org/abs/2301.05062)\ncompiles RASP programs into actual weights. Reusing Tracr was not appealing to\nme for a few reasons. I wanted to express any computational graph, in ordinary\nPython, and when I read the RASP paper I didn’t find the language particularly\nintuitive. I’m a firm believer in the principle articulated by Feynman, “What I\ncannot create, I do not understand.” And to be honest, building it sounded fun.\n\nThat compiler became [torchwright](https://github.com/physicsrob/torchwright).\nYou define a computation graph in ordinary Python; torchwright produces the\nweights of a transformer that executes it. There is no training anywhere in the\npipeline.\n\nOne decision mattered more than any other for making the problem tractable: I\ndid not start with a transformer anyone would recognize from a model card. I\nstarted with the simplest substrate I could get away with — FFNs with ReLU,\nthe easiest nonlinearity to derive constructions in (also the [Attention is All\nYou Need](https://arxiv.org/abs/1706.03762) nonlinearity); a position encoding\nof my own design (the same mechanism as a learned position embedding, just with\nvalues I chose instead of trained); and no normalization layers at all. A decoder-only transformer in\nshape, but with every part chosen for my convenience. Pushing it to a standard\narchitecture came later.\n\n## Where values live\n\nThe first thing the compiler has to manage is memory. The residual stream is the only shared memory there is. Think of it as a whiteboard with numbered columns: every value the graph computes owns a set of columns for as long as anything downstream still needs it. The columns need not be contiguous. The compiler scatters weight matrices to wherever the inputs happen to sit.\n\nA transformer layer comprises two sublayers: attention, followed by a FFN. Each\nsublayer has a skip connection that bypasses it, meaning that the sublayer can\nonly add to the residual stream. It computes `out = in + f(in)`\n\n, and the skip\nconnection rearranges nothing. That one fact drives the whole compilation\nscheme. A new value lands in zeroed columns Once a value is no\nlonger needed by subsequent layers, its associated columns are freed by writing\nits negation When a computational graph calls for two values to\nbe added, we can use the skip connection itself to represent the add.\n\n## Simple obvious operations first\n\nWith the plan in mind, I tackled simple obvious operations first, and built up to increasing complexity.\n\nAdd came first — adding two values is trivially a linear operation. In fact the substrate is embarrassingly redundant about it: an attention head can write the sum into fresh columns, or — if one of the inputs is no longer needed — move one value into the other’s columns and let the skip connection do the arithmetic. The FFN can play either role too. That is four ways to write add, before even trying.\n\nComparison came next. `equals_vector(x, c)`\n\nasks whether an embedding-valued\ninput matches a constant vector assuming has the same magnitude as\nToken embeddings are exactly the vectors this wants — normalized, roughly\northogonal — so “which token am I looking at?” becomes a few `equals_vector`\n\ncalls.\n\nThe formulation here is also fairly straightforward:\n\nworks out to at a match, and for anything whose dot product falls at least short of is a sharpness constant; is the margin a mismatch has to clear, and embeddings are constructed so that mismatches clear it.\n\n`select(cond, t, f)`\n\nwas the next operation. Effectively an if/else. I ended up\nadopting the convention of as true and as false, and torchwright is\nconsistent on this point throughout.\n\nAssuming that the boolean `cond`\n\nis and that and are smaller than\nsome big constant (the compiler sizes from the value bounds it propagates\nthrough the graph), select can be computed as\n\nThe trick: shoves the chosen branch safely positive, so it survives the ReLU, and shoves the other branch negative, so the ReLU kills it to zero.\n\nMost of the library (the ReLU version, anyway) is variations on that move: a big constant, a ReLU, and the unpicked branch dying on the wrong side of it.\n\nFrom there the library grew by composition: multi-way switch, boolean logic, table lookups that map an embedding-valued input through a FFN to an embedding-valued output, etc.\n\nSoon I expanded to handle sequence data — this is a transformer, after all.\n\nThe first position scheme was a custom encoding, inspired by the original Attention is All You Need sinusoids, but only using a handful of columns with the others set explicitly to zero. The original scheme adds position to every column — cancellable in principle, since the encodings are known, but I preferred position confined to its own columns, leaving the rest of the whiteboard clean by construction.\n\nThe cross-position primitives — attention heads that read a value from a fixed offset back, or find the most recent value written at a position where some condition held — were built against those columns, and the machine ran on that scheme for months.\n\n## As standard as I could push it\n\nAt that point I had what I set out to build: graphs compiled, ran, and gave the correct answers. But they ran on a somewhat custom transformer, using my own forward-pass implementation. I wanted to see if I could push this all the way to a standard, off-the-shelf, huggingface transformer.\n\n### RMSNorm\n\nA stock model normalizes the residual stream in every layer, with two steps: divide the stream by its root-mean-square, then multiply each column by a per-column gain. By using a dedicated column with a large value, I was able to deliberately peg the RMS to a known large value. The gain is then chosen to undo the normalization, making the RMSNorm layers effectively the identity.\n\n### Position\n\nThe vast majority of modern LLMs use [rotary position\nembeddings](https://arxiv.org/abs/2104.09864) (RoPE), so naturally I\nwondered whether it was a fit for torchwright. The good news was that in RoPE\nnothing gets added to the residual stream — which means none of our\ncalculations get corrupted. Position is totally independent of the residual\nstream. Instead of adding to the residual stream, RoPE treats pairs of\ndimensions as objects, and rotates these paired dimensions by an angle\nproportional to their position. Many architectures use *partial\nRoPE* — rotating only some of the dimensions — a convention I adopted.\n\nMy position-aware primitives all had to be rebuilt on top of rotation, and each turned out to have a clean RoPE-native construction:\n\n**Position-independent matching.** Put the content on the unrotated dimensions.**Reading a fixed offset back.** The key and query are chosen as constant vectors, where all dimensions are the same, except the key values are prerotated so that the rotation induced by RoPE cancels out at exactly one position, a fixed offset in the past. Everywhere else the different frequency components add less coherently and are suppressed.**The current absolute position.** RoPE deliberately hides absolute position, but a head that attends to the beginning-of-sequence token leaks it: the softmax weight landing on that one token shrinks predictably as the sequence grows, and a piecewise-linear layer inverts that weight back into the integer position. (The inversion stays monotone out past 60,000 positions, with worst-case error about a third of the half-integer threshold that matters for rounding.)\n\n### The nonlinearity\n\nOriginally I had derived every construction in ReLU because that’s what the\noriginal Attention is All You Need used, and it seemed like the easiest form to\nderive operations. But modern models use a gated FFN called\n[SwiGLU](https://arxiv.org/abs/2002.05202), and in the back of my head I had\nhoped I would eventually be able to target the Llama architecture. A strong\ngoal of mine was to eventually deliver a compiled artifact (a checkpoint) that\nruns on vanilla huggingface with not a single line of custom code.\n\nAs it turns out, Llama was more difficult for my use case because it doesn’t support partial RoPE, but stock Phi-3 supported everything I needed, as long as I switched to SwiGLU. My constructions were all derived in ReLU, but a Phi-3 FFN is gated: instead of linear -> activation -> linear, its hidden layer is a product, where and is the sigmoid.\n\nFortunately a lot of my ReLU work was usable even in SwiGLU, thanks to the observation that sharpens the swish function to the point where it’s almost indistinguishable from a ReLU. You just fold the factors of 128 into the weight matrices on either side, and it’s effectively a ReLU.\n\nMultiplication actually comes out far better than under ReLU: exactly, for all and because The gate’s smoothness cancels instead of approximating. Under ReLU, multiplying two values took a table-lookup construction; under Swish it is trivial.\n\nSome constructions came out a lot cleaner than their ReLU originals. Select’s big-offset apparatus dies entirely under the gate: a complementary pair,\n\nselects exactly — no offset constant, no bound on the branches at all.\n\nI ended up keeping both ReLU and SwiGLU implementations around. You can import\nthe primitives from either `torchwright.ops.relu`\n\nor `torchwright.ops.swiglu`\n\n— whichever architecture you want to target.\n\nWith normalization, position, and the nonlinearity replaced, the compiled models\nstopped being mine in any architectural sense: they are ordinary `transformers`\n\ncheckpoints in the Phi-3 architecture — causal softmax attention, RoPE, RMSNorm,\ngated-SiLU FFNs, a KV cache — and `AutoModelForCausalLM`\n\nloads them with no\ncustom code and no `trust_remote_code`\n\n.\n\n## Bounded approximations\n\nSome of the constructions are exact, but plenty are piecewise-linear approximations. Correctness gets measured rather than assumed. Every approximated operation is measured against its exact-math reference, and the resulting error bounds are committed to the repository. Graph nodes can carry assert predicates. During development each assert is checked against the exact reference evaluation and again during debug forward passes of the compiled model. One that passes in exact math but fires when compiled points at exactly the operation where approximation error exceeded its budget. And when things get really confusing, there is a probe that diffs a compiled transformer node-by-node against direct evaluation of the source graph.\n\nOne practical note: the compiled models run in fp32. bf16 won’t work, but consider what bf16 is — roughly two decimal significant figures, a rough approximation that trained models tolerate because they were trained with it. Nothing here was trained, so nothing here learned to shrug off that kind of noise.\n\n## What it can do\n\nOne thing about the execution model, because it is easy to mis-picture: a compiled graph computes exactly one thing — the next token, given the previous tokens in the sequence. Ordinary autoregressive generation loops it: run the graph, append the emitted token, run it again. Some graphs lean on that loop, keeping their algorithm’s state in the tokens they have already emitted rather than anywhere inside the network. Loops whose bounds are known at compile time can instead unroll into depth, the way the adders’ carry chains do. At the output end, the graph’s result is embedding-valued; the tied embedding table turns it into logits, and greedy decoding picks the token.\n\nThe repository contains twelve example graphs: three different adders, a Caesar cipher whose shift amount is a runtime input, digit sorting, a two-level loop, Fibonacci, a family of four calculators that all compute the same functions in very different ways — and the binary increment from the top of the post. Compiling that one is a single function call, from a checkout of the repository:\n\n``` python\nfrom examples.binary_increment import create_network_parts, D_MODEL, D_HEAD\nfrom torchwright import compile_hf_bundle\n\noutput_node, embedding = create_network_parts()\ncompile_hf_bundle(output_node, embedding, \"binary_increment_hf_bundle\",\n                  d=D_MODEL, d_head=D_HEAD, optimize=1)\n```\n\nThat call lands the graph in a 16-layer decoder at hidden size 512, in under\nten seconds on a laptop CPU. The layer count comes from the compiler’s\nscheduler, which assigns every operation to a layer: independent operations can\nshare a layer, but an operation that reads another’s output has to wait for a\nlater one. Depth is therefore the graph’s critical path, not its size — those\n16 layers are essentially the length of the increment’s parse-carry-emit chain.\n(What depth costs as algorithms grow is the next post’s subject.) The result is\nexactly the checkpoint from the top of the post: it loads with a stock\n`pipeline()`\n\ncall, increments binary strings, and contains not one trained\nweight.\n\n## What’s next\n\nWith the tool in place, I could finally go back to the question that started all of this: what arithmetic is a transformer actually capable of expressing? I ended up compiling four calculators, each computing expressions the same way from the outside:\n\n```\nprint(generate(\"12*12\\n\", return_full_text=False)[0][\"generated_text\"])\n# 144\n```\n\nInside, they could hardly be more different. One does multiplication the way you learned it in school. One works like a hardware multiplier. One writes its scratch work out as tokens and reads it back. And one does no arithmetic at all — it memorized every answer — yet for small enough numbers it is the shallowest of the four. What each approach costs, in layers, in parameters, and in tokens, is tackled in the next post.\n\n## References\n\n- Pérez, Barceló & Marinković (2021).\n[Attention is Turing-Complete](https://jmlr.org/papers/v22/20-302.html).*JMLR*22(75). - Merrill & Sabharwal (2023).\n[The Expressive Power of Transformers with Chain of Thought](https://arxiv.org/abs/2310.07923). ICLR 2024. - Abdin et al. (2024).\n[Phi-3 Technical Report](https://arxiv.org/abs/2404.14219). - Weiss, Goldberg & Yahav (2021).\n[Thinking Like Transformers](https://arxiv.org/abs/2106.06981)(RASP). ICML 2021. - Lindner et al. (2023).\n[Tracr: Compiled Transformers as a Laboratory for Interpretability](https://arxiv.org/abs/2301.05062). NeurIPS 2023. - Vaswani et al. (2017).\n[Attention Is All You Need](https://arxiv.org/abs/1706.03762). NeurIPS 2017. - Su et al. (2021).\n[RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864). - Wang & Komatsuzaki (2021).\n[GPT-J-6B](https://github.com/kingoflolz/mesh-transformer-jax)— the origin of partial rotary (rotating only some dimensions). - Shazeer (2020).\n[GLU Variants Improve Transformer](https://arxiv.org/abs/2002.05202)(SwiGLU).\n\n## Citation\n\nRobert Porter. \"Introducing torchwright.\" Out of Distribution, July 2026. https://ood.dev/posts/torchwright-intro/\n\n```\n@misc{porter2026torchwrightintro,\n  author       = {Porter, Robert},\n  title        = {Introducing torchwright},\n  year         = {2026},\n  month        = {jul},\n  howpublished = {\\url{https://ood.dev/posts/torchwright-intro/}},\n  note         = {Out of Distribution (blog)}\n}\n```\n\n", "url": "https://wpnews.pro/news/torchwright-compile-computation-graphs-into-vanilla-transformer-weights", "canonical_source": "https://ood.dev/posts/torchwright-intro/", "published_at": "2026-07-24 17:17:06+00:00", "updated_at": "2026-07-24 17:23:02.199441+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "large-language-models", "ai-research", "developer-tools"], "entities": ["Torchwright", "Rob Halsmith", "Phi-3", "Hugging Face", "GitHub", "RASP", "Tracr"], "alternates": {"html": "https://wpnews.pro/news/torchwright-compile-computation-graphs-into-vanilla-transformer-weights", "markdown": "https://wpnews.pro/news/torchwright-compile-computation-graphs-into-vanilla-transformer-weights.md", "text": "https://wpnews.pro/news/torchwright-compile-computation-graphs-into-vanilla-transformer-weights.txt", "jsonld": "https://wpnews.pro/news/torchwright-compile-computation-graphs-into-vanilla-transformer-weights.jsonld"}}