nn.nu

nn.nu — neural-network layers on the grad autograd tape.

Every layer here is a pure TAPE BUILDER: it takes a * GTape and input GVars (activations and pre-registered weight/const GVars) and returns a GVar, recording ordinary grad ops. Nothing here owns parameters or touches the device — the caller registers weights with grad_param / grad_const, calls backward(loss) once, and grad derives every gradient. Device replay (gput) and megakernel fusion (gpfuse) therefore come free, because an nn graph is just a grad graph.

The math is lifted verbatim from two places that had proven it and were drifting apart: grad's PyTorch-float64 LoRA-block oracle (grad/tests/lora_block.nu + lora_oracle.sh, loss and all 14 adapter gradients to 1e-11) and nurllama's finetune path (__ft_rms / __ft_lora_lin / __ft_rope / __ft_headslice). This package is the single home for that math; the two new layers (nn_layernorm, nn_swiglu) carry their own finite-difference checks.

Convention: 2-D activations are [rows=T, cols]. A weight is [in, out] (x[T,in] · W[in,out]); a bias / norm-weight is a [out] row broadcast over rows. Row reductions (rmsnorm/layernorm mean) use a shared ones-column [n,1] the caller threads in via nn_ones, so a block builds it once.

API

@ nn_const * GTape tp ( Vec f ) v i rows i cols → GVar

A [rows, cols] (or [cols] when rows==0) const tensor from flat data.

@ nn_param * GTape tp ( Vec f ) v i rows i cols → GVar

A [rows, cols] parameter tensor from flat data (registered for grad).

@ nn_ones * GTape tp i n → GVar

A ones column [n,1] — the row-reduction operand for the norms.

@ nn_linear * GTape tp GVar x GVar w → GVar

x[T,in] · W[in,out].

@ nn_linear_bias * GTape tp GVar x GVar w GVar b → GVar

x[T,in] · W[in,out] + b[out] (bias broadcast over rows).

@ nn_lora_linear * GTape tp GVar x GVar w0 GVar a GVar b f scale → GVar

LoRA linear: x·W0 + scale·(x·A)·B, with A[in,r], B[r,out], scale = α/r. W0 is typically a frozen const; A/B are the trainable params.

@ nn_rmsnorm * GTape tp GVar x GVar w GVar ones i h f eps → GVar

RMSNorm: x ⊙ rsqrt(mean(x²)+eps) ⊙ w. ones is a [h,1] column; the row mean is x²·ones / h.

@ nn_layernorm * GTape tp GVar x GVar w GVar b GVar ones i h f eps → GVar

LayerNorm: (x - mean) / sqrt(var + eps) ⊙ w + b, per row over h features. mean = x·ones/h; var = mean((x-mean)²) = (x-mean)²·ones/h. w is [h] (scale), b is [h] (shift); ones is [h,1].

@ nn_silu * GTape tp GVar x → GVar

SiLU / swish: x · sigmoid(x).

@ nn_swiglu * GTape tp GVar gate GVar up → GVar

SwiGLU: SiLU(gate) ⊙ up — the gated-MLP activation (gate, up are the two projections of the same input).

@ nn_softmax * GTape tp GVar x i axis → GVar

Softmax over axis — named alias of grad's op for symmetry.

@ nn_rope_fill i t i hd f base ( Vec f ) cos_out ( Vec f ) sin_out → v

Fill NEOX/half-split rope tables cos/sin for t positions × hd/2 frequencies, θj = base^(-2j/hd). Host-side; the caller const-ifies the two [t, hd/2] vectors (nnconst ... t half).

@ nn_rope * GTape tp GVar h GVar cosc GVar sinc i t i hd → GVar

NEOX half-split rope on one [t, hd] head: split into halves x1,x2 and rotate — [x1·cos - x2·sin, x2·cos + x1·sin]. cosc/sinc are [t, hd/2].

@ nn_head * GTape tp GVar q i t i hi i hd → GVar

Slice head hi of a [t, nhd] projection: columns [hihd, hi*hd+hd).

@ nn_causal_mask_fill i t ( Vec f ) out → v

Fill a causal mask [t, t]: 0 on/below the diagonal, -1e9 above.

@ nn_gqa_attention * GTape tp GVar q GVar k GVar v GVar cosc GVar sinc GVar mask i t i nh i nkv i hd f iscale → GVar

Grouped-query attention over pre-projected q[t,nhhd], k/v[t,nkvhd]. RoPE is applied to each q/k head; nkv key/value heads are shared across nh query heads (kv index = qi / (nh/nkv)). mask is [t,t]; iscale is the 1/sqrt(hd) score scale. Returns the concatenated context [t, nh*hd].

@ nn_cross_entropy * GTape tp GVar logits GVar onehot GVar onesV → GVar

Cross-entropy for next-token / classification over [rows, V] logits with a one-hot target [rows, V]. picked = (softmax(logits) ⊙ onehot)·onesV is the target probability per row; CE = -mean(log(picked)). onesV is a [V,1] column. (Rows whose onehot is all-zero contribute log(0); slice or mask those out upstream — see nn_cross_entropy_rows.)

@ nn_cross_entropy_rows * GTape tp GVar logits GVar onehot GVar onesV i keep → GVar

Cross-entropy over the FIRST keep rows only (next-token training leaves the last position without a target; keep = rows-1). Slices the picked column to [keep,1] before the log-mean.