The Tensor Computer

A complete von Neumann computer — ALU, registers, cache, memory, disk, 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, RAM with an NTM-style content-addressed region, a 4KB-sector hard disk, a GPU, and peripheral ports — implemented entirely as differentiable tensor operations in JAX. (Two-level virtual memory with a TLB and a softmax-arbitrated system bus exist too, as standalone modules; the step function does not call them yet.) 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. One cycle of the soft machine, step by step:

One cycle of the tensor computer's step() A soft program counter fetches one 112-wide instruction from the learned program by dot product and it is decoded and sharpened at temperature τ; a fault cascade (PC, opcode, register) decides how much of the cycle commits. Operands are soft register reads. The ALU computes all 32 operations and blends them by opcode weight, setting flags; memory is read and written softly through a router into the L1 cache, RAM and the NTM; push, pop, call and return go through the stack; the GPU draws into the framebuffer and six peripheral ports transfer. Writeback mixes the ALU result, the loaded value, the popped value and the I/O transfer into rd, and the next PC blends branch target, return address and PC + 1. Every update is x + w·(x_taken − x). program I_mem 2^A × 112, learned PC soft one-hot next PC branch · return · +1 fetch / decode PCᵀ · I_mem, sharpened at τ fault cascade PC · opcode · register L1 cache learned eviction registers 16 × 32, soft reads stack push · pop · call · return RAM 64 MB ALU all 32 ops, blended writeback into rd NTM 256 slots memory router soft read + write peripherals 6 ports flags from the blended result GPU plot · fill · blit framebuffer 480 × 640 × 3 next PC = blend of branch target, return address, PC + 1 rd = mix of ALU result, loaded value, popped value, I/O gated by LOAD / STORE weight decides how much of the cycle commits soft I/O, every cycle Every unit computes every cycle. Each update is x + w·(x_taken − x), with w the weight of the instruction being that kind of instruction.
One cycle of step(): every unit computes every cycle, and each update is blended by how much the current instruction is that kind of instruction.

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 — float32 throughout, jax.nn.softmax for all soft addressing, jax.grad for end-to-end differentiation, JIT-compiled and vectorized across a batch of programs.
  • Architecture: 16×32-bit registers, a direct-mapped L1 cache with a learned eviction network, 64MB RAM plus a 256-slot NTM-style region with content-based addressing, a disk window with 4KB sectors, a 64K-word tensor GPU, and a 480×640×3 display framebuffer — each a tensor in the machine-state tuple. Two-level page tables with a TLB and the system bus are implemented as separate modules outside the step.
  • 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 add through bubble_sort, binary_search, and gcd.

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 gradient through the annealed softmax stack is badly behaved. On one four-instruction kernel its norm is 2e-5 at τ = 10, about 1 at τ = 1, and 9e-6 at τ = 0.01: five orders of magnitude, non-monotone, and in a different place for a different program. 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, these are the results the original paper reports. The August 2026 re-implementation described on this page has not reproduced them: its own status file says that nothing in it yet synthesizes a correct non-trivial program, because the shipped training budgets are deliberately toy-sized and the learned kernels are wrong at them. Treat the table as the paper's claim, not this code's:

TaskEpisodes (as reported)Learned program
add847 ± 123 to 100%ADD R0,R0,R1; HALT — optimal
max2,341 ± 412 to 100%CMP; JGE; MOV; HALT — see below
sum_array15Kcorrect loop, generalizes to any n
bubble_sort45K to 90%, with a curriculumnested-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 programs the paper reports 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, and the clip below is that target environment. The point of building the screen into the machine is that it is a literal H×W×3 float32 framebuffer inside the machine state, so an agent's actions and what it sees could live in one differentiable graph. That is the aim, not something this clip proves: the windowing system that actually runs today, tensor-os, runs on the discrete VMs, which share the soft machine's ISA and framebuffer layout. The pixel-art look is deliberate because it keeps the rendering pipeline inspectable.

What the video is actually showing

The GUI-agent target environment. The machine's display is an H×W×3 float32 tensor in its state; the GUI that runs on it today does so on the discrete VMs, not inside the JAX graph.

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:

The paper: the full machine-state tuple, the soft-addressing formalization, the five-phase annealing schedule, and the two-week research program.

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.