grad.nu

grad — reverse-mode automatic differentiation over tensor.

Define-by-run: every g_* op computes its value EAGERLY through the tensor package and appends one node to a flat TAPE. backward(loss) walks the tape once in reverse — node inputs always have smaller indices, so reverse index order IS topological order — accumulating dL/dx into a gradient slot per node. Fan-out (one value used twice) sums naturally.

Ownership: the tape is a SINGLE-OWNER ARENA. It owns every value tensor and every gradient tensor; grad_param/grad_const COPY the caller's tensor in, gvar_value/grad_of hand out BORROWS (aliased Tensor views — never free them, invalid after tape_free / tape_reset_to), and one tape_free releases everything. No per-node ownership, no per-node closures — a ( Vec GNode ) of {op, a, b, scalar} plus two parallel pointer vecs. This keeps the hot loop free of NURL's closure-capture hazards and gives a deterministic reverse order for free.

The minibatch pattern: register parameters once, take tape_mark, then per batch build the episode's graph, backward, step the optimizer, and tape_reset_to(mark) — parameters (and their Adam state, which lives in the optimizer) survive, intermediates are freed, the arena keeps its capacity.

M1 scope (CPU): elementwise ops + all-axes reductions. Binary ops require EQUAL shapes (broadcast backward is M2, with its own oracle); scalar variants g_adds/g_muls cover the mixed case. A shape mismatch POISONS the tape (tape_ok → F, ops keep returning without touching memory) rather than half-computing.

Verification: tests/grad_test.nu checks every backward rule two independent ways — central finite differences and hand-derived analytic identities (e.g. d/dx Σx² = 2x, exact to the bit).

API

@ gop_param → i

── op codes ──────────────────────────────────────────────────────────

@ gop_const → i

@ gop_add → i

@ gop_sub → i

@ gop_mul → i

@ gop_div → i

@ gop_neg → i

@ gop_adds → i

@ gop_muls → i

@ gop_relu → i

@ gop_sigmoid → i

@ gop_tanh → i

@ gop_exp → i

@ gop_log → i

@ gop_sqrt → i

@ gop_sum → i

@ gop_mean → i

@ gop_matmul → i

@ gop_bmm → i

@ gop_transpose → i

@ gop_reshape → i

@ gop_softmax → i

@ gop_slice → i

@ gop_concat → i

: GNode

: GNode {
    i op
    i a
    i b
    f s
}

One tape node: the op, its input node ids (-1 = none), a scalar operand.

: GVar

: GVar {
    i id
}

A variable on the tape — an opaque node handle.

: GTape

: GTape {
    i ok  // 1 healthy · 0 poisoned (shape mismatch etc.)
    ( Vec GNode ) nodes
    ( Vec s ) vals  // *Tensor per node (tape-owned)
    ( Vec s ) grads  // *Tensor per node, 0 until backward touches it
    ( Vec s ) aux  // *GAux per node, 0 for ops that need none (slice starts)
}

: GAux

: GAux {
    ( Vec i ) starts
}

Per-node auxiliary payload (only g_slice uses one today).

@ tape_new → *GTape

@ _g_tfree s pp → v

@ _g_auxfree s pp → v

@ tape_free * GTape tp → v

@ tape_ok * GTape tp → b

@ tape_len * GTape tp → i

@ tape_mark * GTape tp → i

Watermark for tape_reset_to: everything appended after mark is dropped.

@ tape_reset_to * GTape tp i mark → v

Drop every node at id >= mark (freeing its tensors) and ZERO the gradients of what remains — a fresh episode over the surviving parameters. The vecs keep their capacity, so a minibatch loop does not re-malloc the arena.

@ _g_heap Tensor t → s

Move a Tensor VALUE into a fresh heap cell the tape owns. The value's Vec handles alias into the cell — the caller must NOT free the value after.

@ _g_val * GTape tp i id → s

@ _g_grad_ptr * GTape tp i id → s

@ _g_poison * GTape tp s why → GVar

@ _g_push * GTape tp i op i a i b f sc Tensor val → GVar

Append a node whose value is val (ownership moves to the tape).

@ _g_same_shape s pa s pb → b

Equal-shape check for M1 binops.

@ grad_param * GTape tp Tensor w → GVar

Register a PARAMETER (requires-grad leaf). The tensor is COPIED in; the live, optimizer-updated copy is the tape's (read it via gvar_value).

@ grad_const * GTape tp Tensor c → GVar

Register a CONSTANT (no gradient flows into it).

@ gvar_value * GTape tp GVar v → Tensor

The node's value as a BORROWED Tensor view (do not free; invalid after tape_free / a reset past this node).

@ grad_of * GTape tp GVar v → Tensor

The accumulated gradient as a BORROWED Tensor view. Allocates a zero gradient on first touch so the borrow is always valid.

@ g_scalar * GTape tp GVar v → f

First element of the node's value — the scalar-loss readout.

@ _g_view s pp → Tensor

A borrowed Tensor view of a tape value (alias — never free).

@ _g_binop * GTape tp GVar a GVar b i op → GVar

@ g_add * GTape tp GVar a GVar b → GVar

@ g_sub * GTape tp GVar a GVar b → GVar

@ g_mul * GTape tp GVar a GVar b → GVar

@ g_div * GTape tp GVar a GVar b → GVar

@ _g_unary * GTape tp GVar a i op f sc → GVar

@ g_neg * GTape tp GVar a → GVar

@ g_adds * GTape tp GVar a f sc → GVar

@ g_muls * GTape tp GVar a f sc → GVar

@ g_relu * GTape tp GVar a → GVar

@ g_sigmoid * GTape tp GVar a → GVar

@ g_tanh * GTape tp GVar a → GVar

@ g_exp * GTape tp GVar a → GVar

@ g_log * GTape tp GVar a → GVar

@ g_sqrt * GTape tp GVar a → GVar

@ _g_reduce * GTape tp GVar a i op → GVar

@ g_sum * GTape tp GVar a → GVar

@ g_mean * GTape tp GVar a → GVar

@ g_mse * GTape tp GVar y GVar t → GVar

Mean squared error as a composite: mean((y − t)²). Exercises fan-out-free chaining; the FD harness covers it end to end.

@ g_matmul * GTape tp GVar a GVar b → GVar

@ g_bmm * GTape tp GVar a GVar b → GVar

@ g_transpose * GTape tp GVar a → GVar

@ g_reshape * GTape tp GVar a ( Vec i ) shape → GVar

shape is borrowed (copied here; tensor_reshape adopts the copy).

@ g_softmax * GTape tp GVar a i axis → GVar

@ g_slice * GTape tp GVar a ( Vec i ) starts ( Vec i ) stops → GVar

starts/stops are borrowed; starts is retained (aux) for the backward scatter.

@ g_concat * GTape tp GVar a GVar b i axis → GVar

@ _g_acc_reduce * GTape tp i dst Tensor c f sgn → v

Accumulate an out-shaped contribution c into node dst's gradient, summing over broadcast axes (numpy alignment: trailing dims line up, a size-1 or missing dst dim receives the sum over that out axis). sgn is +1/-1 (sub/div's second input negates). Row-major walk of c — one documented, deterministic accumulation order.

@ _g_ensure_grad * GTape tp i id → v

Allocate node id's gradient as zeros (same shape as its value) if absent.

@ backward * GTape tp GVar loss → b

dL/d(node) of loss seeds to ones; every earlier node receives the sum of its consumers' contributions. Returns F on a poisoned/invalid tape.

---

gput.nu

grad/gput.nu — the GPU replay engine: run a captured tape episode on the device, bit-exactly.

The CPU tape (grad.nu) stays the semantic reference. gput_capture takes ONE recorded episode — parameters, constants, every op node up to the loss — and mirrors it onto the device: a value buffer and a gradient buffer per node, plus the stride/offset metadata each kernel needs. After capture the graph is STATIC (the mark/reset training loop re-records the same node structure every episode anyway): per minibatch the caller uploads fresh input rows into their const slots (gput_set_input), replays forward and backward on the device, and steps the device optimizer — nothing but the loss scalar needs to come back to the host.

Bit-exactness discipline (the aegpu recipe, generalized per op):

• Every kernel reproduces the CPU implementation's floating-point

ROUNDING AND ORDER: accumulation chains are spelled with explicit __dadd_rn/__dsub_rn/__dmul_rn/__ddiv_rn (never fused — NVRTC's default fmad contraction would change results), inner products and reductions run serially in the CPU's documented index order inside one thread, and broadcast-reduce accumulation walks each slot's contributions in the same row-major subsequence order as grad.nu's gacc_reduce.

• Optimizer scalars that need libm pow (Adam's lr_t) are computed on the

HOST with the same expressions opt.nu uses and shipped to the kernel.

• exp/log-based unaries (sigmoid, tanh, exp, log, softmax) mirror the

CPU FORMULAS exactly, but the device's exp()/log() are CUDA libm: on the gpu package's CPU backend they are glibc and the results are bit-identical; on real CUDA hardware they can differ by a few ulp. Everything else — relu, +, −, ×, ÷, sqrt, matmul/bmm, sum/mean, data movement, SGD/Adam — is bit-exact on BOTH backends. An MLP/AE workload (relu + linear + mse + Adam) is therefore bit-exact end-to-end everywhere; tests/gput_parity_test.nu pins both tiers.

Coverage: every grad.nu op. Restrictions (capture poisons with a message): TE_F64 tapes only, and g_bmm only with both operands carrying the full batch (no batch broadcast).

Ownership: GProg owns every device buffer; one gput_free. The tape is borrowed during capture and NOT needed afterwards (except as the target of gput_param_sync_host, which writes trained parameter values back into the tape's own tensors so gvar_value keeps working).

API

@ _gp_src → s

── the kernels ─────────────────────────────────────────────────────── One source, compiled once per kit (cached by entry name). All f64.

: GpNode

: GpNode {
    i op
    i a
    i b
    f s
    i n  // value element count
    i reach  // 1 = loss ancestor (backward touches it)
    ( Vec i ) dims
    GkBuf val
    GkBuf grad
    GkBuf scr  // out-shaped scratch (mul/div backward)
    GkBuf meta  // i64 stride/offset blocks
}

One mirrored node. dims caches per-op host-side launch geometry: binop broadcast: [offB1 offB2 offB3 offBA offBB tfreeA tfreeB] matmul [M K N] · bmm [B M K N] · transpose [M N] softmax/concat [outer axlen inner (da)]

: GProg

: GProg {
    b ok
    * GpuKit kit
    ( Vec GpNode ) nodes
    i loss
}

: GpOpt

: GpOpt {
    b ok
    i kind  // 0 sgd · 1 adam
    f lr
    f clip
    i t
    ( Vec i ) ids
    ( Vec f ) alphas
    ( Vec GkBuf ) m
    ( Vec GkBuf ) v
    GkBuf ptab  // [dptr len]* per param (built lazily for the clip norm)
    GkBuf nrm  // 1-elem norm result
}

Device optimizer over a captured program's parameters — opt.nu mirrored: per-param L2, global-norm clip, the same runtime 1−β Adam arithmetic, the step counter behind the heap pointer.

@ _gp_nobuf → GkBuf

@ _gp_bfree GkBuf b → v

@ _gp_push_vec ( Vec i ) dst ( Vec i ) src → v

Push nd entries of v (aligned to nd out dims, missing lead = filler).

@ _gp_accred_block ( Vec i ) meta ( Vec i ) oshape ( Vec i ) ost ( Vec i ) eff * u tfree → v

The acc-reduce block for one input: [nd ost dec fdim fdec], plus tfree. dec = row-major strides of the input's dims aligned into out dims (1 on broadcast dims); fdim/fdec describe the odometer over the broadcast dims.

@ _gp_ew_block ( Vec i ) meta ( Vec i ) ost ( Vec i ) xe ( Vec i ) ye → v

An ew block [nd ost xe ye].

@ _gp_fail * GProg pg s why → v

@ _gp_upload_tensor * GpuKit kit * Tensor t → GkBuf

Upload a CPU tensor's data into a fresh f64 device buffer.

@ gput_capture * GpuKit kit * GTape tp GVar loss → *GProg

Mirror one recorded episode ([0, tape_len)) onto the device. loss is the node backward seeds from. Check gput_ok before use.

@ gput_ok * GProg pg → b

@ gput_free * GProg pg → v

@ _gp_node * GProg pg i id → GpNode

@ gput_set_input * GProg pg GVar v ( Vec f ) data → b

Fresh values for an input slot (a const's minibatch rows, or a param).

@ _gp_run * GProg pg s name i grid i block ( Vec i ) args → b

@ _gp_fill_buf * GProg pg GkBuf b f v → b

@ _gp_fwd_node * GProg pg i k → b

One forward launch for node k. Leaves are data — nothing to do.

@ gput_forward * GProg pg → b

Recompute every non-leaf value on the device, in tape order.

@ _gp_accred * GProg pg GpNode dst GpNode nd GkBuf src i moff i tfree f sgn → b

The out-shaped contribution src accumulated into input dst's gradient over the broadcast axes (block at moff, odometer size tfree).

@ _gp_scr_ew * GProg pg GpNode nd GkBuf x GkBuf y i moff i op → b

g (out-shaped) ⊙ other-input → scr, via the ew block at moff.

@ _gp_bwd_node * GProg pg i k → b

@ gput_backward * GProg pg → b

Zero every gradient, seed dL/d(loss) = 1, sweep the tape in reverse over the loss-ancestor set, then zero const-leaf gradients (grad.nu's epilogue).

@ gput_value * GProg pg GVar v ( Vec f ) out → b

Download a node's value into out (resized by the caller to node size).

@ gput_grad * GProg pg GVar v ( Vec f ) out → b

@ gput_loss * GProg pg → f

The loss scalar (node value element 0).

@ gput_param_sync_host * GProg pg * GTape tp → b

Write every parameter node's device value back into the CPU tape's own tensor (in place), so gvar_value reflects the trained weights.

@ gpopt_new i kind f lr → *GpOpt

@ gpopt_sgd_new f lr → *GpOpt

@ gpopt_adam_new f lr → *GpOpt

@ gpopt_free * GpOpt o → v

@ gpopt_set_clip * GpOpt o f maxn → v

@ gpopt_add * GpOpt o * GProg pg GVar p f alpha → v

Register one parameter (its Adam moments start at zero, on the device).

@ gpopt_step * GpOpt o * GProg pg → b

One update step from the gradients on the device — opt.nu's semantics: untouched params (not loss ancestors) are SKIPPED, the clip norm covers every registered ancestor's gradient in registration order, Adam's lr_t comes from host pow exactly like the CPU path.

---

opt.nu

grad/opt.nu — optimizers over tape parameters: SGD and Adam, per-parameter L2 (weight decay), optional global-norm gradient clipping.

An optimizer is built empty (opt_adam_new lr / opt_sgd_new lr) and parameters are added one by one with their OWN L2 coefficient (opt_add o tp p alpha) — weights typically carry alpha, biases 0, which is the param-groups need in miniature. opt_step o tp reads each parameter's accumulated gradient from the tape and updates the parameter value IN PLACE on the tape (the tape owns the live copy; read it back with gvar_value).

Semantics: g = grad + alpha·w (L2 applied to the raw gradient; the loss is expected to be mean-scaled — no hidden batch divide here) SGD : w -= lr·g Adam : m = β1·m + (1−β1)·g ; v = β2·v + (1−β2)·g² w -= lr·√(1−β2ᵗ)/(1−β1ᵗ) · m/(√v + ε) β1 = 0.9, β2 = 0.999, ε = 1e-8 — sklearn/PyTorch defaults.

The step counter lives behind the *Opt heap pointer, so it ADVANCES — the frozen-Adam-t bug class (a scalar counter field on a by-value struct never persists in NURL) cannot recur, and tests/opt pin the trajectory bit-for- bit against a hand-computed two-step Adam.

Clipping (opt_set_clip o maxn): one global L2 norm over every added parameter's gradient; if it exceeds maxn every gradient is scaled by maxn/norm before the update — the standard clip-by-global-norm.

A parameter whose gradient was never touched by backward() is SKIPPED (PyTorch's None-grad rule), not decayed.

API

: Opt

: Opt {
    i kind  // 0 sgd · 1 adam
    f lr
    f clip  // 0 = off
    i t  // Adam step count (advances — behind the heap pointer)
    i total  // total moment elements
    ( Vec i ) ids  // param node ids
    ( Vec f ) alphas  // per-param L2
    ( Vec i ) offs  // each param's slab offset into m / v
    ( Vec i ) lens  // each param's element count
    ( Vec f ) m  // first moments (adam)
    ( Vec f ) v  // second moments (adam)
}

@ _opt_new i kind f lr → *Opt

@ opt_sgd_new f lr → *Opt

@ opt_adam_new f lr → *Opt

@ opt_free * Opt o → v

@ opt_set_clip * Opt o f maxn → v

@ opt_add * Opt o * GTape tp GVar p f alpha → v

Register one tape parameter with its L2 coefficient (0 for biases).

@ _opt_gnorm * Opt o * GTape tp → f

The global L2 norm of every registered parameter's gradient (0-grads skip).

@ opt_step * Opt o * GTape tp → v

One update step from the gradients currently on the tape.