Browse documentation
aner.nn · Reference

Differentiable sequence primitives

Compose native embeddings, causal attention and custom objectives with bounded tensor shapes, tested gradients and explicit CPU or Metal/MPS execution.

Implemented in development

Aner provides differentiable tensor composition, embeddings, normalization, token loss and attention in the existing aner.tensor and aner.nn imports. The implementations use native C++ and the existing reverse mode graph. They let a researcher build and modify a decoder from ordinary functions and classes.

These primitives underpin the native generative workflow, including byte/BPE and supported GPT 2 tokenization, training, fine tuning, LoRA, checkpoint/resume, Metal Float32 execution and memory estimates. Current bounds still target small experiments; larger model budgets, arbitrary external architectures and distributed sequence training remain future work. The service interface in aner.llm remains independent.

Run a composed decoder

The complete sequence example implements a custom decoder block with two heads, learned token and position embeddings, causal and padding masks, LayerNorm, GELU, residuals, a feed forward expansion, tied input/output embeddings and named Adam slots.

sh
aner check examples/neural_sequence.aner
aner examples/neural_sequence.aner --prod
aner examples/neural_sequence.aner --dev reports/sequence-first

Use a new report directory. The example trains on a fixed synthetic sequence fixture to check composition and updates. Its loss reduction demonstrates that the native computation and gradients connect; it is not a language quality benchmark. Replace its data and architecture deliberately for research applications, keeping independent training and evaluation partitions.

Shapes and ownership

A differentiable Tensor has positive row major dimensions with rank 1 through 8. Existing tensor literals and two argument constructors still create matrices. Use tensor.reshape to declare batch, time and head axes. Shapes are runtime checked; the language type remains Tensor.

aner
import aner.tensor
import aner.nn

let values = [[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]]
let batch = tensor.reshape(values, shape: List(2, 2, 2))
print(tensor.rank(batch))       // 3
print(tensor.dim(batch, 0))     // 2
let heads = tensor.permute(batch, axes: List(1, 0, 2))

tensor.rank(value) returns rank and tensor.dim(value, axis) returns an extent. Axes are zero based and nonnegative. nn.rows/tensor.rows flatten all leading dimensions, and nn.cols/tensor.cols report the last extent. The legacy two index value access uses that flattened matrix view.

Composition operations preserve immutable contiguous logical values. CPU operations materialize copies; Metal reshape aliases immutable resident storage while permutation/gather operations materialize device results. They do not introduce writable or strided aliases. Copying a Tensor handle shares its immutable node; reshape and permutation create new nodes linked to their source. nn.parameter creates an independent trainable leaf, while nn.detach creates detached data. Gradients, precision/device changes and optimizer updates preserve logical shapes. Equal flattened sizes do not make different logical shapes compatible.

Tensor composition API

CallContract
tensor.reshape(value, shape)Positive List of Int64 extents, same element count. No inferred -1 extent.
tensor.permute(value, axes)List every axis exactly once; forward permutes and backward reverses the mapping.
tensor.slice(value, axis, start, length)Nonempty in bounds interval; backward accumulates into the selected source region.
tensor.concatenate(a, b, axis)Equal rank and matching other axes; gradients reach both inputs.
tensor.broadcast_to(value, shape)Right align dimensions; each source extent must equal its target or be one. Backward sums repeated uses.
tensor.sum(value, axis)Sum one axis and retain it with extent one.
tensor.batched_matmul(a, b)Multiply final two axes and broadcast leading batch axes from the right. Both operands have rank at least two.

For example, [batch,time,width] multiplied by [width,output] yields [batch,time,output]; the projection receives gradients summed across batches and time. [batch,heads,time,width] multiplied by [batch,heads,width,keys] yields [batch,heads,time,keys].

Elementwise arithmetic requires identical logical shapes. Make broadcasts explicit before nn.add or nn.multiply. Pointwise activations, scaling, bias addition and softmax preserve logical shape. Softmax acts on the last axis. nn.argmax replaces the final extent with one. nn.matmul and nn.transpose remain matrix APIs and reject other ranks; use batched products and permutations instead.

Neural primitives

CallContract
nn.embedding(table, ids, shape)Table is [vocabulary,width]. ids is an exact Int64 list whose count matches shape. Output is shape + [width]; IDs must be in range. Repeated IDs scatter add their gradients.
nn.layer_norm(value, gain, bias, epsilon)Normalize the final axis using population variance. Gain and bias are [width] or [1,width] and may be parameters. Epsilon is finite and positive.
nn.gelu(value)Exact error function GELU with its analytic derivative.
nn.dropout(value, probability, seed, training)Inverted dropout with probability in [0,1). A call owns a deterministic SplitMix64 stream; the same seed reproduces the mask. Evaluation and zero probability return the input.
nn.cross_entropy_ids(logits, targets, ignore_index)Stable log sum exp categorical loss. Integer targets flatten the leading logits axes. Ignored rows have zero gradient; average by nonignored target count. A fully ignored batch is an error.
nn.masked_softmax(scores, mask, causal, offset)Last axes are query and key positions; normalize permitted keys only.
nn.attention(query, key, value, mask, causal, offset)Scaled dot product attention composed from batched products, permutation, scaling and masked softmax, with gradients for Q, K and V.

Shape/ID/target/axis lists reject non Int64 elements; token IDs never round through a floating point tensor. Existing nn.cross_entropy remains available for dense one hot targets.

Dropout's training argument is model behavior. --dev and --prod do not change it. Use an explicit seed schedule across training calls rather than unintentionally repeating one mask. No operation here consumes a shared global random stream.

Attention masks and positions

Use Q and K shaped [...,queries,head_width] and [...,keys,head_width], with V shaped [...,keys,value_width]. A typical self attention shape is [batch,heads,time,head_width]. The scores are Q × transpose(K) / sqrt(head_width). V may broadcast into the Q/K batch shape but cannot introduce additional heads or batches.

Masks are detached finite 0/1 tensors, in the same dtype and device as scores. A one permits a key; a zero excludes it. They broadcast from the right to the complete score shape. [batch,1,1,keys] represents key padding across all heads and query positions. [[1.0]] permits every key before applying causality. Masked probabilities and gradients are exactly zero.

With causal: true, key index must be at most query_index + offset. Offset is nonnegative and key positions start at zero. For self attention use offset zero. For queries representing a tail starting at position s of a longer key sequence, pass s. Offset controls masking only; it does not apply positional embeddings or maintain a KV cache. A row with no permitted keys fails explicitly. Pad queries can attend valid keys and then be excluded from the loss using ignore_index.

Masking occurs inside stable softmax without creating infinite tensor entries. To inspect selected probabilities, compose the scores yourself and call nn.masked_softmax; nn.attention does not hide projection weights or learnable parameters.

Custom model and parameter contract

Use ordinary functions or a class with Tensor fields and a forward method. Its output must be a connected Tensor built from differentiable operations. Initialize trainable leaves with nn.parameter; plain tensors remain constants. The decoder example provides the complete class pattern.

Assign a stable name to every distinct trainable parameter when calling nn.update(optimizer, parameter, key). The existing optimizer registers a slot on its first update and checks the complete logical shape, dtype and device on later updates. Explicit class update methods enumerate parameters; automatic parameter discovery is not implemented. An update returns a new leaf and must be assigned back to the field.

For tied weights, use the same parameter handle in both forward paths and update it once. Backward accumulates both uses into that parameter. Replacing each tied field with separate nn.parameter calls breaks that identity. A custom forward can freeze a parameter by using detached data and excluding it from updates. The reference aner.gen decoder already supports full fine tuning, exact name freezing and LoRA; custom classes retain responsibility for their own training and state lifecycle. See the from scratch guide.

nn.backward(loss) replaces reachable gradients. It does not accumulate across calls. Build a fresh forward graph each step, update leaves, and release previous graph handles. Do not retain every historical loss tensor in a list. The tests verify stable tracked buffer use over repeated updates.

Devices precision and resource limits

Composition and sequence kernels support CPU Float64/Float32 and Metal Float32. Unsupported GPU calls fail explicitly without a hidden CPU fallback. Existing dense CUDA support remains available for its implemented operations. CPU reference kernels accumulate in Float64 and round Float32 outputs at tensor boundaries; Metal uses native Float32 arithmetic and reductions. CPU/device comparisons use declared numerical tolerances, while strict future token causality and same backend resume are checked separately.

Bounds remain 1,000,000 elements per tensor, rank at most 8, graph depth 256, at most 4,096 graph nodes and 4,000,000 graph elements. Batched matrix products allow at most 50,000,000 multiply adds per call. The attention implementation materializes scores and probabilities; its workspace grows quadratically with sequence length. These are checked resource limits, not estimates of total process memory. Raising them is not a substitute for implementing scalable memory management and accelerator kernels.

Inspection and validation

Development reports display logical shapes. Their JSON records shape in addition to flattened rows and cols, both for ordinary tensor captures and differentiation nodes. Raw values retain the existing opt in limits. Production uses identical mathematical operations, masks and seeds, without recording inspection. Expensive explicit diagnostic expressions still evaluate unless guarded as described in development mode.

aner.sequence checks Float64 finite differences for composition, broadcasts, batched products, embeddings, LayerNorm, GELU, dropout, token loss and causal Q/K/V attention. It checks repeated IDs, padding, score scaling, position offsets, future token isolation in forward and backward, shape/type errors, Float32 propagation, optimizer shape validation and graph lifetime. aner.sequence_language checks named arguments, custom decoder training, persistent sessions, source diagnostics, logical inspection shapes and dev/prod equality.

Aner documentation · Reference and guides