A fully concurrent, thread-safe autograd engine, written the way Rust wants it written.
poorgrad begins from
Karpathy's micrograd and then
takes the road the others don't: no Rc<RefCell<...>>, no single-threaded
assumption, no graph rebuilt on every pass, and a payload generic over
scalars and tensors alike. Sharing a computation graph across threads is
not a feature bolted on with locks; it is what the types guarantee.
Most autograd engines are define-by-run: the graph is built dynamically as
code executes, mutated in place, and single-threaded by assumption.
poorgrad goes the other way, and everything below follows from that one
choice:
- Record once, run anywhere. Expressions record a static tape;
forwardreplays an O(1) snapshot of it, andbackwardreplays the evaluation's own copy, so runs never lock the graph and never disturb each other. Declared inputs take per-run payloads throughforward_with— feeds are run state, not graph state — so one shared network serves any number of threads, each feeding its own data and differentiating its own target. - Values are
Copy. AValueis a borrow of its network plus a position: operators never consume their operands, handles cross threads freely, and a value outliving its graph is a compile error, not a bug report. Every type'sSend + Synccontract is asserted at compile time. - Mutation is a state transition. A gradient step produces the next network generation in O(parameters): the parameter store is rebuilt, the recorded structure is shared untouched through an append-only arena, replaced payloads are reclaimed with their generation, and older generations stay fully usable. Snapshot isolation, for networks.
- Performance falls out of structure. One
Mutexin the engine itself, taken briefly per operation — the arena insidecow_vecholds the only other, and training never touches it; O(1) forks; nounsafein this crate, with#![forbid(unsafe_code)]keeping it a promise rather than a claim (the arena'sunsafecore iscow_vec's, encapsulated behind its tested interface). CPU-only, on purpose: the engine is the point — and the claims are measured, not asserted:cargo benchruns the suite.
use poorgrad::Network;
let network = Network::new();
let w = network.parameter(0.0_f64);
let x = network.input(0.0);
let y = network.input(0.0);
// Operators record the graph; values are `Copy` and never consumed.
let error = w * x - y;
let loss = error * error;
let w_symbol = w.symbol();
let x_symbol = x.symbol();
let y_symbol = y.symbol();
let loss_symbol = loss.symbol();
// The graph is recorded once; every step feeds one sample of the line
// `y = 2 * x` and steps to the next generation, which shares the recorded
// graph while replacing the parameter payloads.
let samples = [(1.0, 2.0), (2.0, 4.0), (3.0, 6.0)];
let mut network = network;
for step in 0..100 {
let (sample_x, sample_y) = samples[step % samples.len()];
let loss = network.resolve(loss_symbol);
let evaluation = network.forward_with([(x_symbol, sample_x), (y_symbol, sample_y)]);
let gradients = evaluation.backward(loss);
network = network.update(&gradients, |w, g| w - 0.02 * g);
}
let learned = network.resolve(w_symbol).payload().unwrap();
assert!((learned - 2.0).abs() < 1e-6);The threaded example goes further: per-sample gradients computed on separate rayon threads over one shared network, then three learning rates trained in parallel on O(1) forks.
The engine builds, evaluates, differentiates, and trains computation
graphs over a generic payload — scalars or tensors alike. The complete
machinery is grouped into private payload,
engine, and neural modules while the
crate root keeps the public API flat. From tape to training:
Value— aCopyproxy to a value allocated in aNetwork, and the only graph handle in the public API. It borrows the network, so it cannot outlive it. Arithmetic operators build the graph (let x = v1 + v2;allocates a new computed node on the same network) and never consume their operands.Network— the single owner of the state of every value of a graph, backed by the arena-basedcow_veccrate: allocation is append-only, cloning forks the network in O(1), and the whole structure isSend + Sync. A gradient step is a state transition:updateproduces the next generation, rebuilding only the parameter store while sharing everything else.inputdeclares a per-run input with a default payload;forward_withbinds fed payloads to inputs for one run, validated against their recorded shapes.Symbol— a detached,Copyidentifier for a value.Network::resolveturns it into a proxy in a compatible generation, while rejecting unrelated or divergent networks. Training loops keep symbols of the loss and parameters acrossupdatesteps.EvaluationandGradients— the per-run results offorwardandbackward, read back with the sameValueproxies that built the graph. Runs never mutate the network, so any number of them can execute concurrently.Field— a value-aligned buffer tied to a network lineage rather than one generation, with elementwise algebra (+,scale,zip,map).Gradientsis an alias for it: the field one backward run produces, combined across runs and carried across generations as optimizer state (momentum, Adam) with no conversion;updatetakes a compatible field covering the current graph as its update direction.Tensor— the built-in tensor payload: an immutable runtime shape over a shared element buffer read through a strided layout, sotransposeand the broadcasts are O(1) views rather than copies (tensors are immutable, so aliasing a buffer is always safe). Its storage is an extensible representation — a denseArc-shared buffer or a non-allocating constant today, with room for more — and elements are read withiter,as_slice, orto_vec. ANetwork<Tensor<f64>>uses the same graph, evaluation, differentiation, and update APIs as a scalar network. TheTensorialtrait providesmatmul,transpose, the reductionssumandsum_along, and the explicit broadcastsbroadcast_likeandbroadcast_along(scalars implement scalar semantics for the same trait bound). Broadcasting is explicit by design: a single value spread across a named reference's shape, or a payload repeated along one named axis — never an implicit alignment rule. Shapes are inferred and checked when expressions are recorded — a shape mismatch panics at the offending line, before anything runs: the record-once answer to type-level shape checking.Tape— internal: the append-only record (a Wengert list) shared by a network and all of its proxies, and the engine's single synchronization point.Function— internal: a statically sized enum of the differentiable operations, each variant owning its operand links and parameters and implementing theOperationtrait (forward math and gradient routing per operation, dispatched with a plainmatch).Neuron— a scalar-granularity affine unit with weights, a bias, and anActivation. Its parameters are allocated on the network and retained as symbols across compatible generations.Layer— a dense layer at tensor granularity:activation(x.matmul(w) + b)over a[batch, inputs]value, one weight matrix and one bias vector, with the bias explicitly broadcast over the batch axis.Mlp— dense layers described by a sequence of widths such as[3, 4, 4, 1]: tanh hidden layers, an affine output, and caller-controlled initialization from each parameter's requested shape.cross_entropy— the classification loss as a composed formula over the fused, numerically stableValue::log_softmax: the mean negative log-likelihood of one-hot (or soft) targets, fed per run like any other input.
A poor man's autograd: no GPU required, none wanted. The name is the only modest thing about the design.
The vocabulary used across code and docs — the scientific meaning of each term and its mapping to the Rust types — is collected in TERMINOLOGY.md.
chain— build a small expression graph by chainingValueproxies with arithmetic operators, then evaluate it and compute gradients:cargo run --example chain.gradient_descent— fitw * x + bto a line, threaded with rayon: one shared network differentiated for per-sample targets on separate threads, then trained in parallel on O(1) forks, one per learning rate:cargo run --example gradient_descent.mlp_xor— train a tanh MLP on XOR at tensor granularity: the graph is recorded once, and every training step feeds a different minibatch throughforward_withwhile the tape never grows:cargo run --example mlp_xor.
Licensed under either of MIT or Apache-2.0 at your option.
