The Tensor Computer
A complete von Neumann computer — ALU, registers, cache, virtual memory, bus, GPU, peripherals — rebuilt as differentiable tensor operations in JAX, so that programs written in machine code can, in principle, be learned end-to-end by gradient descent.
Problem
An LLM agent that "increments a loop counter" spends 20–50 tokens doing
what one machine instruction — ADD R1, R1, 1 — does in four bytes. That
100–1000× representational overhead compounds across every step of every
task: a 50-step workflow becomes thousands of tokens of inference, when
the equivalent compiled program would finish in microseconds. Worse, the
architecture is a mismatch — transformer autoregression is not isomorphic
with the symbolic composition and recurrence that structured cognition
actually needs.
Mechanistic interpretability hints that the overhead is unnecessary: grokked transformers implement clean, sparse algorithms (Nanda et al. showed modular-addition models computing a literal discrete Fourier transform), and specific behaviors in large models live in sub-1000-dimensional subspaces. If the useful computation is algorithmic, why carry 8GB of weights to run it? The question this project asks: can a program be learned directly as machine code, with gradients, instead of generated as text?
Solution
The Tensor Computer is a complete von Neumann architecture — ALU, register file, flags, cache, NTM-style main memory, a 4KB-sector hard disk, two-level virtual memory with a TLB, a softmax-arbitrated system bus, a GPU, and peripheral ring buffers — implemented entirely as differentiable tensor operations in JAX. Every discrete decision a real CPU makes is replaced by a temperature-scaled softmax:
- The opcode is a distribution over 32 operations; the ALU computes all 32 in parallel and blends them by weight.
- The program counter is a probability distribution over instruction
addresses; instruction fetch is
(PC)ᵀ · I_mem. - Register selectors, memory addresses, and branch conditions are all soft.
At temperature τ → 0 every softmax collapses to one-hot and the machine recovers exact discrete semantics. At higher τ, gradients flow through the entire computation, so a program — encoded as 112 real-valued parameters per instruction — is just another differentiable object you can optimize.
To prove the architecture is real and not a toy, the implementation
carries four independent backends — the JAX machine, a reference C VM, a
Rust VM, and a batched C++ simulator — held in agreement by 59 shared
conformance vectors. Recursive factorial(5) returns 120 in exactly 49
cycles on all four, which is the load-bearing check: the soft machine
and the real machine have to be the same machine.
Two things run on top of it. A self-hosting C compiler: tcc, written
in C, compiled to tensor machine code, occupies 40,942 of the machine's
65,536 instruction slots and — running on the tensor computer — compiles
a C program to output byte-identical to the host-built compiler's. And
tensor-os, a preemptive multitasking kernel with a windowing GUI
written in C: damage-tracking compositor, widget toolkit, bitmap font, and
five applications, in 20,256 slots.
How
- Language / runtime: Python, JAX —
float32throughout,jax.nn.softmaxfor all soft addressing,jax.gradfor end-to-end differentiation, JIT-compiled and vectorized across a batch of programs. - Architecture: 16×32-bit registers, 256-line cache with soft LRU eviction, 256×32 NTM-style main memory with content-based addressing, 64MB disk with 4KB sector loading, two-level page tables + TLB, a 64K-word tensor GPU, and a 480×640×3 display framebuffer — each a tensor in the machine-state tuple.
- Training strategy: five-phase temperature annealing (warm-up → anneal → crystallize → discrete local search → extraction), curriculum learning over program length and task difficulty, shaped rewards plus auxiliary losses for credit assignment, hierarchical subroutine libraries, and population-based training for solution diversity.
- Validation harness: a self-hosting C compiler emitting tensor
machine code; an algorithmic task suite from
addthroughbubble_sort,binary_search, andgcd.
Tests
The honest result first: the Tensor Computer never learned through SGD on its own. The soft loss landscape at medium temperature does not track the discrete one — a blend of 32 operations has different semantics from any single operation — and the gradients through the annealed softmax stack are too sharp to descend. Pure gradient-based program synthesis stalled.
What did work was scaffolding. With temperature annealing, a length/difficulty curriculum, shaped rewards, imitation warmstart, and discrete local search all stacked on top of each other, the system learned correct machine-code programs for basic tasks. The caveat is worth stating plainly: with that much scaffolding the "learner" is doing very little learning on its own — the search procedure carries most of the weight. It is closer to guided program search than to an autonomous gradient learner.
Within those bounds the results are real and verifiable:
| Task | Episodes to 100% | Learned program |
|---|---|---|
add | 847 ± 123 | ADD R0,R0,R1; HALT — optimal |
max | 2,341 ± 412 | CMP; JGE; MOV; HALT — see below |
sum_array | 15K | correct loop, generalizes to any n |
bubble_sort | 50K | nested-loop bubble sort, 98.5% @ n=4 |
The max row needs a correction I only found while building the full
implementation, and it is more interesting than the original claim. That
four-instruction sequence is not a signed maximum at full width. This
ISA derives its conditional branches from the zero and negative flags
alone — there is no N != V form, no branch on overflow, and no
instruction that can read the overflow flag into a register — so
CMP a, b; JGE decides a >= b only while a - b cannot overflow. It is
wrong on roughly a quarter of 32-bit input pairs: max(0x7FFFFFFF, -1)
returns -1, halting cleanly with a plausible answer.
The overflow-safe form costs eight instructions, not four. The same defect turned up independently in three places — the evaluation suite's reference implementations, the C compiler's own code generator, and the paper — which is the sort of thing you only catch by making several implementations agree on shared test vectors rather than testing each against itself.
Fixed high temperature (no annealing) plateaus at ~10% — chance. Annealing is not optional; it is the experiment.
Results
The trained programs are correct machine code, and some are genuinely
alien. Asked for absolute value, one run rediscovered the branchless
two's-complement trick (SRA to build a sign mask, then XOR/SUB) —
a technique most programmers would not reach for. An in-place swap was
solved with the classic XOR swap, rediscovered from scratch.
sum_array trained on multiple input sizes converged on a real
loop that length-generalizes to arrays it never saw; trained on a single
size, it overfit into an unrolled straight-line solution — a clean,
legible picture of generalization failure.
The longer-horizon goal is GUI agents, but the important detail in the
clip below is not the UI polish. The "screen" is a literal
H×W×3 float32 tensor inside the machine state, rendered by JAX as part
of the Tensor Computer's primitive video pipeline. The cursor is not a
browser cursor moving over a mockup; it is another peripheral state being
advanced by a program running inside the differentiable computer graph.
In other words, the demo is showing a full computer running inside a tensor graph: CPU state, memory, GPU/framebuffer, peripherals, display scanout, and mouse movement are all tensors updated by the same compiled JAX step function. The pixel-art look is deliberate because it keeps the rendering pipeline inspectable while the research target remains much deeper than UI aesthetics.
What the video is actually showing
A differentiable von Neumann computer in JAX, with a full display
framebuffer represented as an H×W×3 float32 tensor and
a cursor rendered through the machine's own primitive video path.
Full architecture, mathematical formalization, training strategy, and the $1,000 / two-week research program are in the paper, which is here rather than behind a link — the state-tuple definitions and the annealing schedule are the part of this project that is hardest to take on trust, and they are on pages you can read without leaving:
If you want the short version of what to look for: the machine-state tuple is the honest measure of how complete the architecture is, and the annealing schedule is where the negative result lives.
Lessons
The central lesson is about where the difficulty actually lives. Building a differentiable computer is mostly bookkeeping — every component is a tensor, every dispatch is a softmax, and JAX makes the whole thing differentiable for free. The hard part is that differentiable and learnable are not the same thing. A gradient exists at every point and still points nowhere useful, because the soft relaxation of a discrete program is a different function than the program itself, and the relaxation gap is exactly where the optimizer gets lost.
So the project lands as an honest negative-leaning result: gradient descent did not learn programs here, and the working system is better described as heavily-scaffolded program search than as a learner. That is still worth having — it sharpens the question of when gradients help program synthesis, and the differentiable substrate remains a clean testbed for interpretability and for the RL-driven phases of the research program.
The second lesson came from the implementation rather than the research, and I did not expect it. Building four independent backends and forcing them to agree on shared test vectors found bugs that no single implementation could have found, because each one was internally consistent and passed its own tests. All four had independently chosen a different random number generator. A timer off-by-one appeared in two of them from the same root cause, arrived at separately. The straight-through estimator was inverted in a way that returned the wrong value and exactly zero gradient — training would have annealed into crystallization and silently stopped learning while still producing plausible output. And the signed-comparison defect above sat in three places at once.
None of those fail loudly. They produce plausible answers, which is
precisely why differential testing across implementations catches them and
unit tests do not. For a project whose central claim is that two
implementations of the same machine — one soft, one discrete — must agree
to within 1e-4, that turned out to be the methodology rather than an
accident of it.
Neighborhood