tang — differentiable computing in Rust

* Forward + backward, seq 128, half-precision tiles vs. PyTorch bf16 autocast (0.376 vs. 1.329 ms), RTX 3090. Full 4-layer step: 5.05 ms vs. PyTorch 4.40.

How it works

Trace

fn f<S: Scalar>(
    x: S,
) -> S {
    x * x.sin()
}

Plain Rust, generic over the number type

Transform

A graph you can differentiate, fuse and tile

Lower

Native code for whatever chip you have

One pipeline. Rust traces to a graph; the graph is differentiated, fused and tiled; the result lowers to native code for each backend.

f(x, y, z) =

+ Gradient

loading…

Try it. The whole pipeline, compiled to WASM and running in this tab. Edit the expression; the kernel and its gradient regenerate as you type.

Same idea, every scale

The same three steps run from a single dual number to an eight-billion-parameter model. Nothing below uses a separate code path.

A scalar

Dual<f64> f(x) = x · sin(x)

x = 0.00

val = 0.00

dual = 0.00

A scalar. Dual<f64>: value and exact slope in one pass.

A robot arm

Mode / score

–

Shoulder / elbow

–

Hand error

–

Catch. Move the pointer over the panel to steer the hand of rose, our humanoid: a 227 mm upper arm and 236 mm forearm with joint limits. You choose where the hand goes; tang solves the joint angles every frame by damped least squares, with the Jacobian from two dual passes through the forward kinematics. Left alone, it runs a pick-and-place loop.

An orbit

Status

–

Your best miss

–

tang's miss

–

Beat the optimizer. Drag from the launch pad to fling a probe at the target: three shots, gravity does the rest. Then let tang try. It differentiates the miss distance through 1,600 leapfrog steps of the same gravity code with Dual<f64> and walks speed and heading downhill.

A model

This is actually Rust code running in your browser!

tang-train The Quantum Poet 21K params · character-level · physics haikus

// build the model
let mut model = Sequential::<f64>::new(vec![
    Box::new(Linear::new(240, 64, 42)),
    Box::new(Tanh::new()),
    Box::new(Linear::new(64, 64, 137)),
    Box::new(Tanh::new()),
    Box::new(Linear::new(64, 30, 256)),
]);

// train: Trainer takes the model, an optimizer and a loss
fn loss(p: &Tensor<f64>, t: &Tensor<f64>) -> (f64, Tensor<f64>) {
    (cross_entropy_loss(p, t), cross_entropy_loss_grad(p, t))
}

let losses = Trainer::new(&mut model, ModuleAdam::new(0.005), loss)
    .epochs(200)
    .fit(&mut loader);

Click Train to start…

Seed: Temp: 0.8

quantum_poet.rs → More examples →

A model. Train a network in this tab, in Rust compiled to WASM.

A transformer

7.6×

Derived flash attention vs. the hand kernel, n = 1024

+18%

Whole transformer block, derived vs. hand-written, seq 512

0

Lines of op-specific code in the compiler

RMSNorm, RoPE, GQA attention, SwiGLU: forward and backward, all derived tensor-ir.md →

A transformer. Kernels derived from the IR beat hand-written ones.

An LLM

tang-llm RTX 3090 Speed of light: 220 tok/s

Decode, device step122 tok/s · 55% SOL

Decode, 4k context104 tok/s · 54% SOL

Prefill 512, f16 tensor cores774 tok/s

8.2 msDevice step per token

506Kernels per decode step, 470 lowered by the IR

52–217 tok/sResident server with speculative drafting, by acceptance rate

535–757 tok/sServer prefill

1.7 × 10⁻⁴Max logit error vs. an f32 PyTorch reference

64 / 64Greedy tokens matching the reference

+ speculative decoding · MoE offload · iOS · Vulkan on Adreno

An LLM. Qwen3-8B, 4-bit, batch 1, on an RTX 3090.

What you can build

Each of these is an example in the repository. Pick one, paste the command, and it runs in about a minute.

Seven starting points. Each command runs from a fresh clone of the repository.

How it compares

PyTorch is the reference for machine learning; candle and burn are strong Rust frameworks. tang covers more ground: math, physics and ML on one differentiable core.

tangPyTorchcandleburn

Kernels derived from one IRSix backends from one tensor IRtorch.compile generates Triton; most ops hand-writtenHand-written per backendCubeCL kernels generated from Rust

GPU backendsCUDA, Metal, Vulkan, WGSL, OpenCL, CCUDA, ROCm, MPS, CPUCUDA, Metal, CPUCUDA, ROCm, Metal, Vulkan, WGPU, CPU

Forward-mode dualsAny S: Scalartorch.func.jvpNot a built-in focusNot a built-in focus

Symbolic graphs + difftang-exprNot a built-in focusNot a built-in focusNot a built-in focus

Geometry & spatial algebraVec/Quat/Mat, exact predicatesNot a built-in focusNot a built-in focusNot a built-in focus

Dense & sparse solversLU, SVD, QR, Eigen, L-BFGS, LMtorch.linalgPartialNot a built-in focus

Constant-time cryptoMod-q tensors, taint-checkedNot a built-in focusNot a built-in focusNot a built-in focus

LLM servingOpenAI-compatible, speculativeVia vLLM and othersExamplesNot a built-in focus

no_std / WASMBothNeitherWASMBoth

Ecosystem maturityYoungEnormousGrowingGrowing

built in partial or via another project not a built-in focus. Hover a mark for detail. Spot something wrong? Open an issue.

How tang works

Four ideas carry the whole library. Each is small enough to state in one equation, and one example, , runs through all of them.

1Dual numbers carry the derivative

Extend every number with an infinitesimal part , where . Taylor expansion then stops after one term, so evaluating a function on a dual number returns its value and its exact derivative together:

(1)

For the running example, ordinary arithmetic on duals gives

(2)

a

value part

ε part, the derivative

Two numbers per number. Drag . Each operation updates the value part as usual and the ε part by its derivative rule: multiplies by , and the product rule adds .

In tang this is just a type. Code written once over S: Scalar computes values with f64 and derivatives with Dual<f64>, with no tape and no symbolic engine. For many inputs, the same functions trace into a graph instead, and tang-expr runs reverse mode over it.

2Two primitives: views and comprehensions

Above the scalar level, every tensor operation is one form. A comprehension maps each output index to a reduction over some inputs, each read through a view, an affine map from indices to memory:

(3)

Matmul, softmax, convolution, gather and attention differ only in their views and their monoid :

(4)

Views are index maps. Hover a cell on the right to see which source element it reads. Transpose, broadcast, slice and stride compose by substitution and never touch memory, so every layout choice stays visible to the compiler.

3Optimizations are derived, not written

Because everything is one form, the compiler's passes are rewrites over that form. Fusion is a single rule: substitute a map into its consumer. Gradients are comprehensions too. The most striking rewrite turns softmax attention into flash attention. Written naively, attention needs the row maximum before the exponentials:

(5)

The stabilize_softmax pass rewrites the max and the two sums into one reduction whose state is updated per key, so the scores never reach memory:

(6)

Same answer, no score matrix. Left, attention as written: every score is stored, then read back for the max and both sums. Right, the derived reduction: keys stream past once while three running values carry everything needed. Memory written per query row is shown below each.

That combine is associative, so it tiles and parallelizes like any sum. The derived kernel runs 7.6× faster than tang's hand-written attention at , and the compiler contains no attention-specific code.

4Lowering to each chip

A scheduled graph prints directly as source for each backend. Here is the running example with its derivative, exactly as the live demo at the top emits it. Each line is one node of the graph, and shared subexpressions are computed once:

(7)

One graph, any backend. Hover a node to find its line, or a line to find its node: each t variable is one node. The same nodes print as CUDA, WGSL, Metal or C; for tensor graphs, schedules add tiling, tensor-core fragments and fused epilogues.

tang.