From 2b5c6165a9590c330db64964c83fbfa32e2a1003 Mon Sep 17 00:00:00 2001 From: nisbenz Date: Mon, 27 Jul 2026 12:26:21 +0100 Subject: [PATCH] add documentation --- README.md | 480 ++++++++++++--- docs/autograd_engine.md | 805 +++++++++++++++++++++++++ docs/decoder_implementation.md | 990 +++++++++++++++++++++++++++++++ docs/neural_network_modules.md | 1010 ++++++++++++++++++++++++++++++++ docs/tensor_mechanics.md | 709 ++++++++++++++++++++++ 5 files changed, 3912 insertions(+), 82 deletions(-) create mode 100644 docs/autograd_engine.md create mode 100644 docs/decoder_implementation.md create mode 100644 docs/neural_network_modules.md create mode 100644 docs/tensor_mechanics.md diff --git a/README.md b/README.md index 7198e23..c478017 100644 --- a/README.md +++ b/README.md @@ -1,129 +1,445 @@ # TensorLib -A from-scratch deep learning tensor library in C with autograd, Transformer modules, and CPU optimizations. - -## Features - -- **Tensor library** — N-dimensional float32 arrays with NumPy-style broadcasting, strided views (transpose, reshape, slice, expand, squeeze), and reference-counted storage -- **Autograd engine** — Dynamic reverse-mode automatic differentiation over 23 differentiable operations, with stale-graph detection and gradient accumulation -- **Neural network modules** — Linear, Embedding, LayerNorm, Dropout, Multi-head Causal Self-Attention, Decoder Block (pre-norm), full Decoder stack, MLP -- **Loss functions** — Cross-entropy (with built-in log-softmax), Softmax, LogSoftmax -- **Optimizers** — SGD, AdamW (decoupled weight decay, bias correction, gradient clipping) -- **Checkpointing** — Versioned, atomic (transaction-safe) save/load with optimizer and RNG state -- **SIMD acceleration** — AVX2+FMA matmul micro-kernel with tile-based blocked algorithm -- **OpenMP parallelism** — Multi-threaded matmul kernels -- **Deterministic RNG** — Splitmix64 PRNG with uniform and normal (Box-Muller) distributions -- **No external ML dependencies** — Pure C, no Python, CUDA, or third-party ML libraries required +A from-scratch deep learning tensor library written entirely in C99 — providing N-dimensional tensors with autograd, Transformer building blocks, loss functions, optimizers, checkpointing, and SIMD-accelerated matmul. **Zero external ML dependencies.** -## Build +> Inspired by [PyTorch](https://github.com/pytorch/pytorch), [ggml](https://github.com/ggerganov/llama.cpp/tree/master/ggml), and [OpenBLAS](https://github.com/OpenMathLib/OpenBLAS). -### CMake (recommended) +--- -```sh -cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -cmake --build build --config Release +## Key Features + +- **Tensor Core** — N-dimensional `float32` arrays with NumPy-style broadcasting, zero-copy strided views, and reference-counted storage +- **Autograd Engine** — Dynamic reverse-mode automatic differentiation over 23 differentiable operations with stale-graph detection +- **Neural Network Modules** — Composable module hierarchy with Linear, Embedding, LayerNorm, Dropout, Multi-head Causal Self-Attention, MLP, and full GPT-style Decoder +- **Loss Functions** — Cross-entropy (numerically stable), Softmax, LogSoftmax +- **Optimizers** — SGD and AdamW (decoupled weight decay, bias correction, gradient clipping) +- **Checkpointing** — Versioned, atomic (transaction-safe) binary save/load with optimizer and RNG state +- **SIMD Matmul** — AVX2+FMA blocked micro-kernel with packed-RHS optimization (MR=4, NR=16) +- **Deterministic RNG** — Splitmix64 PRNG with uniform and normal (Box-Muller) distributions +- **CPU only, no Python, no CUDA, no third-party ML libraries** + +--- + +## Architecture + +```mermaid +graph TB + subgraph Applications + A1["tiny_lm
Byte-level LM"] + A2["mnist_mlp
Digit Classifier"] + A3["autograd_example
Computation Graph"] + end + + subgraph NN["Neural Network Modules"] + N1["Linear, Embedding, LayerNorm"] + N2["Multi-Head Causal Attention"] + N3["Decoder Block (Pre-Norm)"] + N4["Full Decoder Stack"] + N5["Loss Functions"] + N6["Optimizers (SGD, AdamW)"] + N7["Checkpointing"] + end + + subgraph AG["Autograd Engine"] + G1["23 Differentiable Ops"] + G2["Graph Construction"] + G3["Backward Pass"] + G4["Broadcast Gradient Reduction"] + end + + subgraph TC["Tensor Core"] + T1["Allocation & Views"] + T2["Element-wise Ops"] + T3["Reductions"] + T4["Gather"] + end + + subgraph SK["SIMD Kernels"] + S1["AVX2+FMA Matmul"] + S2["Packed RHS"] + S3["Batched Strided"] + end + + A1 --> N4 + A2 --> N1 + A3 --> G1 + + N4 --> N3 + N3 --> N2 + N2 --> N1 + N4 --> N5 + N6 --> G3 + N7 --> G2 + + N1 --> G1 + N2 --> G1 + N5 --> G1 + + G1 --> T2 + G3 --> T2 + G4 --> T3 + + T1 --> T2 + T2 --> T4 + T3 --> T4 + + T4 --> S1 + S1 --> S2 + S1 --> S3 ``` -Requires a compiler with OpenMP support (GCC, Clang, MSVC). +### Layer Summary -### Makefile (GNU Make, GCC) +| Layer | Files | Lines | Purpose | +|-------|-------|-------|---------| +| **Tensor Core** | 7 `.c` + 2 `.h` | ~1,880 | N-dim tensors, broadcasting, views, matmul | +| **Autograd Engine** | 7 `.c` + 2 `.h` | ~1,122 | Reverse-mode AD, 23 ops, graph traversal | +| **NN Modules** | 14 `.c` + 1 `.h` | ~1,837 | Layers, attention, decoder, module system | +| **Losses** | 1 `.c` | ~131 | Cross-entropy, softmax | +| **Optimizers** | 3 `.c` | ~480 | SGD, AdamW, grad clipping | +| **RNG** | 1 `.c` | ~43 | Splitmix64, uniform, normal | +| **Serialization** | 1 `.c` | ~617 | Checkpoint save/load | +| **Total** | 34 files | ~6,110 | | -```sh -make # build and run all tests -make tiny-lm # build the TinyLM example -make mnist # build the MNIST example -make clean -``` +--- -Profile-guided optimization (PGO) for TinyLM: +## Documentation -```sh -make pgo-tiny-lm CORPUS=corpus.txt +Comprehensive documentation is available for each component: + +| Document | Description | +|----------|-------------| +| [**Tensor Mechanics**](docs/tensor_mechanics.md) | Storage model, strided views, broadcasting, AVX2 matmul kernel, API reference | +| [**Autograd Engine**](docs/autograd_engine.md) | Computation graph, 23 differentiable ops, backward pass algorithm, gradient reduction | +| [**Neural Network Modules**](docs/neural_network_modules.md) | Module system, all layers, loss functions, optimizers, checkpointing | +| [**Decoder Implementation**](docs/decoder_implementation.md) | GPT-style decoder stack, multi-head attention, causal masking, training guide | + +--- + +## Component Deep-Dive + +### Tensor Core + +Fixed-precision (`float32`) n-dimensional arrays with: + +- **Strided layout** — element accessed via `offset + sum(coords[i] * strides[i])` +- **Zero-copy views** — reshape, transpose, slice, squeeze, expand share storage +- **Reference counting** — storage freed when last tensor is destroyed +- **Version counter** — detects stale computation graphs in autograd +- **Broadcasting** — right-aligned NumPy-style broadcasting for all element-wise ops + +```mermaid +graph LR + subgraph Storage + S["float* data | refcount: 3 | version: 5"] + end + + T1["tensor A | shape: [2,3]"] --> S + T2["tensor B (view) | shape: [3,2] | offset: 0"] --> S + T3["tensor C (slice) | shape: [1,3] | offset: 3"] --> S ``` -## Project structure +See [tensor_mechanics.md](docs/tensor_mechanics.md) for the full tensor API, broadcasting rules, and AVX2 matmul kernel internals. + +### Autograd Engine + +Dynamic reverse-mode automatic differentiation — the graph is built eagerly during forward execution: + +```mermaid +sequenceDiagram + participant User + participant AG as Autograd + participant Tensor + + User->>AG: ag_matmul(a, b) + AG->>Tensor: Execute t_matmul forward + AG->>AG: Create ag_node with backward fn + AG->>Tensor: Create result ag_tensor (creator = node) + AG-->>User: Return ag_tensor + User->>AG: ag_backward(result) + AG->>AG: Reverse topological sort + loop For each node in reverse order + AG->>AG: Check storage versions (stale detection) + AG->>Tensor: Compute input gradients via backward fn + AG->>AG: Reduce gradients if broadcast + AG->>AG: Accumulate on leaf tensors + end ``` -include/tensorlib/ # Public API headers - tensor.h # Core tensor structs and operations - tensor_matmul.h # Matmul internals - autograd.h # Autograd engine API - nn.h # Neural network module API - -src/ # Source implementation - tensor/ # Tensor ops, views, reductions, matmul - autograd/ # Autograd forward/backward ops, graph traversal - nn/ # Module system, layers (Linear, Embedding, LayerNorm, - # Dropout, MultiheadAttention, DecoderBlock, Decoder, MLP) - losses/ # Cross-entropy, Softmax, LogSoftmax - optim/ # SGD, AdamW, gradient clipping - init/ # PRNG and weight initialization - serialization/ # Checkpoint save/load - -tests/ # Test suite - unit/ - tensor/ # 6 test files - autograd/ # 9 test files - nn/ # 16 test files - optim/ # 2 test files - -examples/ - autograd_example.c # Computation graph demo - tiny_lm/ # Byte-level decoder language model (~1.9M params) - mnsit/ # MNIST MLP classifier - -benchmarks/ - matmul/ # Matmul performance benchmarks + +See [autograd_engine.md](docs/autograd_engine.md) for the full algorithm including broadcast-gradient reduction and transactional error handling. + +### Neural Network Modules + +C-style OOP module system with function-pointer dispatch: + +```mermaid +classDiagram + class nn_module { + +const char* type_name + +char* name + +forward_fn forward + +destroy_fn destroy + +nn_module* parent + +nn_parameter** params + +nn_module** children + +bool training + } + + class nn_linear { + +nn_module base + +ag_tensor* weight + +ag_tensor* bias + } + + class nn_layer_norm { + +nn_module base + +ag_tensor* gamma + +ag_tensor* beta + } + + class nn_multihead_attention { + +nn_module base + +ag_tensor* qkv_weight + +ag_tensor* out_weight + +int n_heads + } + + class nn_decoder { + +nn_module base + +nn_embedding* token_embed + +nn_positional_embedding* pos_embed + +nn_decoder_block** blocks + +nn_layer_norm* final_norm + +nn_linear* lm_head + } + + nn_module <|-- nn_linear + nn_module <|-- nn_layer_norm + nn_module <|-- nn_multihead_attention + nn_module <|-- nn_decoder + nn_decoder o-- nn_decoder_block + nn_decoder_block o-- nn_multihead_attention ``` -## Examples +See [neural_network_modules.md](docs/neural_network_modules.md) for all layers, loss functions, optimizers, and checkpointing. -### Autograd demo +### Decoder Stack -```sh -./build/autograd_example +GPT-2-style causal transformer decoder: + +```mermaid +flowchart TB + Input["Input Token IDs [batch, seq]"] + + TE["Token Embedding: vocab -> d_model"] + PE["Positional Embedding: seq -> d_model"] + Add["+ Addition"] + + B1["Decoder Block 1"] + B2["Decoder Block 2"] + B3["Decoder Block N"] + + LN["Final LayerNorm"] + LH["LM Head (Linear): d_model -> vocab"] + Output["Output Logits [batch, seq, vocab]"] + + Input --> TE + Input --> PE + TE --> Add + PE --> Add + Add --> B1 + B1 --> B2 + B2 --> B3 + B3 --> LN + LN --> LH + LH --> Output ``` -Builds a computation graph `input @ weights + bias -> exp -> mean`, backpropagates, and prints all gradients. +Each decoder block follows the pre-norm pattern (GPT-2): +1. LayerNorm -> Multi-Head Causal Self-Attention -> Residual +2. LayerNorm -> MLP (FFN) -> Residual + +See [decoder_implementation.md](docs/decoder_implementation.md) for multi-head attention internals, causal masking, and the TinyLM training guide. -### MNIST MLP +--- + +## Quick Start + +### Build ```sh -./build/mnist_mlp +# CMake (recommended) +cmake -S . -B build -DCMAKE_BUILD_TYPE=Release +cmake --build build --config Release + +# Makefile +make # build and run all tests +make tiny-lm # build TinyLM example +make mnist # build MNIST example ``` -784→128 ReLU→10 Softmax MLP trained with SGD and cross-entropy loss. +### Train a Language Model -### TinyLM — byte-level language model +```sh +./build/tiny_lm corpus.txt --steps 5000 --generate 200 --prompt "To be" +``` + +4-layer, 192-width, 6-head decoder transformer (~1.9M params) trained with AdamW on raw byte corpora. See [examples/tiny_lm/README.md](examples/tiny_lm/README.md). + +### Train an MNIST Classifier ```sh -./build/tiny_lm corpus.txt --steps 1000 --generate 200 --prompt "Hello" +./build/mnist_mlp ``` -4-layer, 192-width, 6-head decoder transformer (~1.9M params) trained with AdamW on raw byte corpora. See [examples/tiny_lm/README.md](examples/tiny_lm/README.md) for details. +784 -> 128 ReLU -> 10 Softmax MLP trained with SGD and cross-entropy loss. + +--- ## Testing -33 unit test executables covering tensors, autograd, all NN modules, optimizers, loss functions, and checkpointing. +33 unit test executables covering every component: ```sh +# Run all tests via CMake cmake --build build --config Release --target test -# or with Makefile: + +# Or via Makefile make test ``` -## Benchmarks +| Category | Tests | Coverage | +|----------|-------|----------| +| Tensor | 6 | Core, alloc, view, ops, reductions, matmul | +| Autograd | 9 | Core, ops, views, gather, reductions, matmul, backward, integration, public API | +| NN Modules | 16 | RNG, parameters, modules, init, linear, embedding, positional, layer norm, dropout, attention, decoder block, decoder, loss, causal mask, checkpoint, MLP | +| Optimizers | 2 | SGD, AdamW | + +--- + +## Performance + +### Matmul Benchmarks -Matmul benchmarks compare TensorLib's blocked AVX2 kernel against OpenBLAS: +Compare TensorLib's blocked AVX2 kernel against OpenBLAS: ```sh make benchmark-compare ``` -## Design notes +The matmul kernel uses a tile-based blocked algorithm: + +```mermaid +graph TB + subgraph Tiling["Matmul Tiling Strategy"] + direction TB + C["C [M x N] Output"] + A["A [M x K] LHS"] + B["B [K x N] RHS (packed)"] + + subgraph Tiles["Tile Loop: MC=64, NC=64, KC=128"] + T1["Micro-kernel: MR=4, NR=16"] + end + end + + A --> Tiles + B --> Tiles + Tiles --> C +``` + +**Kernel parameters:** +- Micro-kernel: MR=4 rows x NR=16 columns (AVX2: 8 floats x 2 registers) +- Tile sizes: MC=64, NC=64, KC=128 +- RHS packing: pre-pack B panels for cache-friendly access +- SIMD: AVX2+FMA with `_mm256_fmadd_ps` + +See [tensor_mechanics.md](docs/tensor_mechanics.md) for kernel implementation details. + +### Optimization Roadmap + +See [optimizations.md](optimizations.md) for the full optimization roadmap covering: +- SIMD vectorization (AVX-512, ARM NEON) +- Kernel fusion (element-wise ops, softmax, LayerNorm) +- OpenMP parallelism +- FlashAttention +- INT8 quantization + +--- + +## Project Structure + +``` +tensorlib/ +├── README.md # This file +├── optimizations.md # SIMD/parallelism optimization roadmap +├── CMakeLists.txt # CMake build +├── Makefile # GNU Make build (GCC, supports PGO) +│ +├── include/tensorlib/ # Public API headers +│ ├── tensor.h # Core tensor structs and operations +│ ├── tensor_matmul.h # Matmul internals (AVX2 kernel) +│ ├── autograd.h # Autograd engine API +│ ├── autograd_internal.h # Internal autograd helpers +│ └── nn.h # Neural network module API +│ +├── src/ # Source implementation (32 .c files) +│ ├── tensor/ # Tensor ops, views, reductions, matmul +│ ├── autograd/ # Autograd forward/backward, graph traversal +│ ├── nn/ # Module system, layers, decoder +│ ├── losses/ # Cross-entropy, softmax +│ ├── optim/ # SGD, AdamW, grad clipping +│ ├── init/ # PRNG and weight initialization +│ └── serialization/ # Checkpoint save/load +│ +├── tests/ # Test suite (33 executables) +│ ├── fixtures/test_common.h # Custom test framework +│ └── unit/ # Tensor, autograd, nn, optim tests +│ +├── examples/ # Example programs +│ ├── autograd_example.c # Computation graph demo +│ ├── tiny_lm/ # Byte-level decoder LM (~1.9M params) +│ └── mnsit/ # MNIST MLP classifier +│ +├── benchmarks/ # Performance benchmarks +│ └── matmul/ # Matmul vs OpenBLAS +│ +└── docs/ # Detailed documentation + ├── tensor_mechanics.md # Tensor layer reference + ├── autograd_engine.md # Autograd engine reference + ├── neural_network_modules.md # NN modules reference + └── decoder_implementation.md # Decoder implementation reference +``` + +--- + +## Build Requirements + +- **Compiler:** GCC or Clang with C99 + OpenMP support +- **Flags:** `-O3 -march=native -mtune=native -flto -fopenmp -fno-math-errno -funroll-loops -fprefetch-loop-arrays` +- **Tested on:** GCC via MSYS2 UCRT64 on Windows + +--- + +## Design Philosophy + +| Decision | Rationale | +|----------|-----------| +| **C99** | Minimal runtime, direct memory control, portable to embedded systems | +| **No external ML deps** | Self-contained; every algorithm is implemented from scratch | +| **Reference counting** | Deterministic lifetime, no GC pauses, suitable for real-time | +| **Dynamic graph (define-by-run)** | Natural control flow, no tracing step — like PyTorch, unlike TF 1.x | +| **Pre-norm transformer** | Training stability for deep networks (GPT-2 design) | +| **Fused QKV projection** | Better cache locality in attention | +| **Right-aligned broadcasting** | Matches NumPy/PyTorch semantics | +| **Storage version counter** | Detects stale computation graphs between forward and backward | + +--- + +## References & Inspirations -- **Eager execution** — Graph is built dynamically during forward pass; no JIT compilation -- **Manual memory management** — Reference counting for `Storage`, `ag_tensor`, and `ag_node`; no garbage collector -- **Single-threaded API** — Not thread-safe; assumes single-threaded usage -- **CPU only** — No GPU support -- **SIMD scope** — AVX2+FMA only used in matmul; element-wise ops, activations, norm, softmax are scalar (optimization opportunity tracked in `optimizations.md`) +- **[PyTorch](https://github.com/pytorch/pytorch)** — Autograd engine design, nn.Module pattern, storage/stride model +- **[ggml](https://github.com/ggerganov/llama.cpp/tree/master/ggml)** — Tensor struct layout, stride-based operations, kernel dispatch +- **[OpenBLAS](https://github.com/OpenMathLib/OpenBLAS)** — SGemm microkernel tiling strategy, packed-RHS optimization +- **[nanoGPT](https://github.com/karpathy/nanoGPT)** — Decoder architecture, training loop design +- **[NumPy](https://github.com/numpy/numpy)** — Broadcasting semantics, view model diff --git a/docs/autograd_engine.md b/docs/autograd_engine.md new file mode 100644 index 0000000..eab2e97 --- /dev/null +++ b/docs/autograd_engine.md @@ -0,0 +1,805 @@ +# Autograd Engine + +TensorLib's autograd engine provides **dynamic reverse-mode automatic differentiation** over a define-by-run computation graph. Every differentiable operation eagerly executes the forward pass while simultaneously recording a backward node. When `ag_backward` is called on a loss tensor, the engine traverses the recorded graph in reverse topological order, propagating gradients from output to every contributing leaf. + +## Table of Contents + +1. [Overview](#1-overview) +2. [Core Data Structures](#2-core-data-structures) +3. [Computation Graph Construction](#3-computation-graph-construction) +4. [Tensor Lifecycle in Autograd](#4-tensor-lifecycle-in-autograd) +5. [The 23 Differentiable Operations](#5-the-23-differentiable-operations) +6. [Backward Pass (The Core Algorithm)](#6-backward-pass-the-core-algorithm) +7. [Broadcast Gradient Reduction](#7-broadcast-gradient-reduction) +8. [Error Handling & Graph Consistency](#8-error-handling--graph-consistency) +9. [API Reference](#9-api-reference) +10. [Test Coverage](#10-test-coverage) +11. [Design Decisions & Tradeoffs](#11-design-decisions--tradeoffs) + +--- + +## 1. Overview + +TensorLib's autograd system is a from-scratch C99 implementation of the same fundamental idea found in [PyTorch's autograd engine](https://github.com/pytorch/pytorch/blob/main/torch/csrc/autograd/engine.cpp): a **dynamic computation graph** built eagerly during forward execution, traversed in reverse during backward. + +### What autograd provides + +- **Dynamic reverse-mode AD**: the graph is built on-the-fly as `ag_add`, `ag_matmul`, etc. are called. No separate tracing or compilation step. +- **23 differentiable operations** spanning element-wise arithmetic, activations, reductions, views, matrix multiplication, and gathering. +- **Automatic broadcast-reduction**: when inputs are broadcast during forward, gradients are automatically reduced back to input shape during backward. +- **Gradient accumulation**: repeated calls to `ag_backward` accumulate gradients on leaf tensors, matching PyTorch's behavior. +- **Transactional safety**: if any tensor is modified between forward and backward, the entire backward pass is rejected and existing gradients are left unchanged. +- **Manual reference counting**: deterministic lifetime management with no garbage collector. + +### Comparison to PyTorch + +| Aspect | TensorLib autograd | [PyTorch autograd](https://github.com/pytorch/pytorch/blob/main/torch/csrc/autograd/engine.cpp) | +|--------|-------------------|------------------------------------------| +| Graph type | Dynamic (define-by-run) | Dynamic (define-by-run) | +| Backward dispatch | Single-threaded reverse topological sort | Multi-threaded task queue with worker threads | +| Node representation | `ag_node` struct with function pointer | [Node/Edge](https://github.com/pytorch/pytorch/blob/main/torch/csrc/autograd/graph_task.h) with `gradient_edge` | +| Stale detection | Storage version counters | Version counters on variables | +| Memory management | Manual reference counting | Shared pointers (`std::shared_ptr`) | +| Gradient accumulation | Transactional merge after traversal | Per-node accumulation with `AccumulateGrad` | + +### Comparison to ggml + +[ggml](https://github.com/ggerganov/llama.cpp/tree/master/ggml) supports backward passes via `ggml_grad` and `ggml_cgraph`, but builds the entire computation graph statically before execution. TensorLib builds the graph dynamically, which means: + +- No separate "build graph" and "execute graph" phases. +- Control flow (loops, conditionals) naturally produces different graph structures on each invocation. +- Memory is freed incrementally via reference counting rather than bulk-freed after graph execution. + +--- + +## 2. Core Data Structures + +### `ag_op` enum — all 23 differentiable operations + +Defined in `include/tensorlib/autograd.h:27-52`: + +```c +typedef enum { + AG_OP_ADD, // a + b + AG_OP_SUB, // a - b + AG_OP_MUL, // a * b + AG_OP_DIV, // a / b + AG_OP_NEG, // -a + AG_OP_EXP, // e^a + AG_OP_LOG, // ln(a) + AG_OP_POW, // a^exponent + AG_OP_SQRT, // sqrt(a) + AG_OP_RELU, // max(0, a) + AG_OP_SIGMOID, // 1 / (1 + e^-a) + AG_OP_TANH, // tanh(a) + AG_OP_GELU, // 0.5 * a * (1 + tanh(sqrt(2/pi) * (a + 0.044715*a^3))) + AG_OP_MATMUL, // matrix multiply + AG_OP_SUM, // sum along dimension + AG_OP_MEAN, // mean along dimension + AG_OP_MAX, // max along dimension + AG_OP_RESHAPE, // reshape (zero-copy view) + AG_OP_TRANSPOSE, // transpose (zero-copy view) + AG_OP_SLICE, // slice along dimension (zero-copy view) + AG_OP_EXPAND, // expand with broadcasting (zero-copy view) + AG_OP_GATHER_ROWS, // select rows by index + AG_OP_MUL_SCALAR, // tensor * float + AG_OP_DIV_SCALAR // tensor / float +} ag_op; +``` + +### `ag_tensor` — the autograd tensor wrapper + +Defined in `include/tensorlib/autograd.h:67-78`: + +```c +struct ag_tensor { + tensor* value; // the underlying numeric data + tensor* grad; // accumulated gradient (NULL until backward) + + int requires_grad; // 1 = participates in gradient tracking + + ag_node* creator; // the node that produced this tensor (NULL for leaves) + + int ref_count; // manual reference count + + int graph_index; // temporary index used during backward traversal; -1 when idle +}; +``` + +An `ag_tensor` wraps a raw `tensor` (see [tensor_mechanics.md](./tensor_mechanics.md)) and adds gradient tracking. The `creator` pointer forms a linked list from output tensors back to their producing nodes, creating the computation graph. + +### `ag_node` — a graph node + +Defined in `include/tensorlib/autograd.h:81-101`: + +```c +struct ag_node { + ag_op operation; // which op created this node + + int input_count; // number of differentiable inputs (1 or 2) + ag_tensor** inputs; // retained input tensors (owned references) + + ag_tensor* output; // non-owning back-pointer to the result + + ag_backward_fn backward; // local backward function pointer + + void* context; // operation-specific saved state (e.g., scalar value, dim) + void (*free_context)(void*); // destructor for context + + uint64_t* input_versions; // storage versions at forward time + uint64_t output_version; // output storage version at forward time + + int ref_count; // manual reference count +}; +``` + +### Class Diagram + +```mermaid +classDiagram + class ag_tensor { + +tensor* value + +tensor* grad + +int requires_grad + +ag_node* creator + +int ref_count + +int graph_index + } + + class ag_node { + +ag_op operation + +int input_count + +ag_tensor** inputs + +ag_tensor* output + +ag_backward_fn backward + +void* context + +void(*free_context)(void*) + +uint64_t* input_versions + +uint64_t output_version + +int ref_count + } + + class tensor { + +Storage* storage + +int ndim + +int* dims + +int* strides + +int offset + } + + class Storage { + +float* data + +int ref_count + +int size + +uint64_t version + } + + ag_tensor --> tensor : value + ag_tensor --> tensor : grad + ag_tensor --> ag_node : creator + ag_node --> ag_tensor : inputs[] + ag_node --> ag_tensor : output (back-pointer) + tensor --> Storage : storage +``` + +### Comparison to PyTorch's Node/Edge design + +PyTorch uses [`Node`](https://github.com/pytorch/pytorch/blob/main/torch/csrc/autograd/graph_task.h) and `Edge` objects where each edge carries a gradient edge (input index + gradient function). TensorLib simplifies this into a single `ag_node` struct: + +- **PyTorch**: `Node` → `Edge` → `Node` (edges carry metadata about which input slot) +- **TensorLib**: `ag_node.inputs[]` directly holds retained `ag_tensor*` pointers; the input index is implicit (array position) + +TensorLib also avoids PyTorch's `AccumulateGrad` leaf node pattern. Instead, `ag_backward` handles leaf accumulation directly via `merge_persistent_gradients`. + +### Comparison to ggml + +ggml uses `ggml_cgraph` (a static graph) containing `ggml_tensor` nodes and `ggml_grad` metadata. The graph is built in a single pass and executed separately. TensorLib's `ag_node` is closer to PyTorch's dynamic approach: each forward op creates and links its node immediately. + +--- + +## 3. Computation Graph Construction + +### How the graph is built + +Every `ag_*` forward operation follows the same pattern: + +1. Execute the raw `t_*` forward operation to produce an output tensor. +2. Call `ag_make_result`, which checks if any input `requires_grad`. +3. If yes, allocate an `ag_node`, record storage versions for stale detection, retain all inputs, and link the node as the output's `creator`. +4. If no input requires gradients, return the result without a creator (untracked leaf). + +### Forward operation lifecycle + +```mermaid +sequenceDiagram + participant User + participant ag_add as ag_add(a, b) + participant t_add as t_add(a->value, b->value) + participant ag_make as ag_make_result() + + User->>ag_add: ag_add(a, b) + ag_add->>t_add: t_add(a->value, b->value) + t_add-->>ag_add: output tensor + ag_add->>ag_make: ag_make_result(output, AG_OP_ADD, 2, inputs, backward_add) + Note over ag_make: Check requires_grad on all inputs + alt At least one input requires grad + ag_make->>ag_make: Allocate ag_node + ag_make->>ag_make: Record input_versions, output_version + ag_make->>ag_make: Retain all input ag_tensors + ag_make->>ag_make: Link node as result->creator + end + ag_make-->>ag_add: ag_tensor* result + ag_add-->>User: ag_tensor* +``` + +### Concrete example: `ag_add(a, b)` + +From `src/autograd/autograd_ops.c:107-108`: + +```c +ag_tensor* ag_add(const ag_tensor* a, const ag_tensor* b) { + return apply_binary(a, b, AG_OP_ADD, t_add, backward_add); +} +``` + +The `apply_binary` helper (`autograd_ops.c:96-105`) calls `t_add` for the forward, then delegates to `ag_make_result`: + +```c +static ag_tensor* apply_binary(const ag_tensor* a, const ag_tensor* b, + ag_op operation, binary_forward_fn forward, + ag_backward_fn backward) { + tensor* output = forward(a->value, b->value); + ag_tensor* inputs[2] = {(ag_tensor*)a, (ag_tensor*)b}; + return ag_make_result(output, operation, 2, inputs, backward, NULL, NULL); +} +``` + +### No graph → no overhead + +When no input requires gradients, `ag_make_result` returns the result without allocating a node (`autograd_core.c:108-111`): + +```c +if (!requires_grad) { + if (free_context != NULL) free_context(context); + return result; // untracked leaf — no creator, no node +} +``` + +This means mixing tracked and untracked tensors is efficient: only tensors involved in gradient computation pay the graph overhead. + +--- + +## 4. Tensor Lifecycle in Autograd + +### Creating autograd tensors + +**From a raw tensor** — `ag_from_owned_tensor` (`autograd_core.c:6-23`): + +```c +ag_tensor* ag_from_owned_tensor(tensor* value, int requires_grad); +``` + +Takes ownership of `value`. Returns an `ag_tensor` with `grad = NULL`, `creator = NULL`, `ref_count = 1`, and `graph_index = -1`. The `requires_grad` parameter is normalized: any non-zero value becomes `1`. + +**Detaching** — `ag_detach` (`autograd_core.c:25-46`): + +Creates a zero-copy leaf alias that shares the same `Storage` but has `requires_grad = 0`, no grad, and no creator. This is useful for stopping gradient flow through a branch while sharing data. + +### Creating gradient tensors + +`ag_full_like` (`autograd_core.c:153-161`) allocates a new tensor with the same shape as a reference but filled with a constant value. Used internally to create zero-valued gradient buffers for slice and max backward passes. + +### Reference counting + +Both `ag_tensor` and `ag_node` use manual reference counting: + +- **`ag_tensor_retain`** / **`ag_tensor_release`**: increment/decrement the tensor's `ref_count`. When it reaches zero, the creator node is released, `grad` and `value` are freed, and the `ag_tensor` itself is freed. +- **`ag_node_retain`** / **`ag_node_release`**: same pattern for nodes. When a node's refcount reaches zero, it frees its `context`, releases all retained input tensors, and frees itself. + +The graph is held alive through reference counting: `ag_make_result` retains each input tensor in the node. The output tensor owns the node through its `creator` pointer. Releasing the output cascades to releasing the node, which cascades to releasing its inputs. + +### ag_tensor lifecycle + +```mermaid +stateDiagram-v2 + [*] --> Created : ag_from_owned_tensor() + Created --> GraphAttached : ag_make_result() + GraphAttached --> BackwardCalled : ag_backward() + BackwardCalled --> GraphAttached : ag_backward() again (accumulate) + BackwardCalled --> Zeroed : ag_zero_grad() + Zeroed --> GraphAttached : ag_backward() + GraphAttached --> Freed : ag_tensor_release() + BackwardCalled --> Freed : ag_tensor_release() + Zeroed --> Freed : ag_tensor_release() + Created --> Freed : ag_tensor_release() +``` + +--- + +## 5. The 23 Differentiable Operations + +Operations are grouped by category. For each operation, the table shows the forward behavior and the backward gradient formula. + +### Element-wise Binary + +| Operation | Forward | Backward (grad_a, grad_b) | +|-----------|---------|--------------------------| +| `AG_OP_ADD` | `a + b` (broadcast) | `grad_out`, `grad_out` | +| `AG_OP_SUB` | `a - b` (broadcast) | `grad_out`, `-grad_out` | +| `AG_OP_MUL` | `a * b` (broadcast) | `grad_out * b`, `grad_out * a` | +| `AG_OP_DIV` | `a / b` (broadcast) | `grad_out / b`, `-grad_out * a / b²` | + +### Scalar Binary + +| Operation | Forward | Backward (grad_a) | +|-----------|---------|-------------------| +| `AG_OP_MUL_SCALAR` | `a * scalar` | `grad_out * scalar` | +| `AG_OP_DIV_SCALAR` | `a / scalar` | `grad_out / scalar` | + +### Unary + +| Operation | Forward | Backward | +|-----------|---------|----------| +| `AG_OP_NEG` | `-a` | `-grad_out` | +| `AG_OP_EXP` | `e^a` | `grad_out * output` (reuses forward output) | +| `AG_OP_LOG` | `ln(a)` | `grad_out / a` | +| `AG_OP_POW` | `a^x` | Special-cased: x=0 → 0, x=1 → grad_out, else `grad_out * x * a^(x-1)` | +| `AG_OP_SQRT` | `√a` | `grad_out / (2 * output)` (reuses forward output) | +| `AG_OP_RELU` | `max(0, a)` | `grad_out` if a>0, `0` if a≤0, `NaN` if a=NaN | +| `AG_OP_SIGMOID` | `1/(1+e^-a)` | `grad_out * output * (1 - output)` | +| `AG_OP_TANH` | `tanh(a)` | `grad_out * (1 - output²)` | +| `AG_OP_GELU` | tanh approximation | `grad_out * d/dx GELU(x)` | + +### Reductions + +| Operation | Forward | Backward | +|-----------|---------|----------| +| `AG_OP_SUM` | `sum(a, dim)` | Expand upstream to input shape (uniform distribution) | +| `AG_OP_MEAN` | `mean(a, dim)` | Same as sum, but scaled by `1/dim_size` | +| `AG_OP_MAX` | `max(a, dim)` | Upstream distributed equally among tied maxima; 0 elsewhere; NaN input → NaN gradient | + +### Views + +| Operation | Forward | Backward | +|-----------|---------|----------| +| `AG_OP_RESHAPE` | Zero-copy reshape | `reshape(grad_out, original_shape)` | +| `AG_OP_TRANSPOSE` | Zero-copy transpose | `transpose(grad_out, dim0, dim1)` (inverse) | +| `AG_OP_SLICE` | Zero-copy slice | Scatter upstream back into zero-filled input shape | +| `AG_OP_EXPAND` | Zero-copy expand | Clone upstream (broadcast gradient reduced by `ag_backward` later) | + +### Matmul + +| Operation | Forward | Backward | +|-----------|---------|----------| +| `AG_OP_MATMUL` | `A @ B` | `grad_out @ B^T`, `A^T @ grad_out` (handles vector/matrix/batched) | + +### Gather + +| Operation | Forward | Backward | +|-----------|---------|----------| +| `AG_OP_GATHER_ROWS` | Select rows by index | Scatter-add upstream into zero-filled table gradient (supports duplicate indices) | + +--- + +## 6. Backward Pass (The Core Algorithm) + +The backward pass is implemented in `src/autograd/autograd_backward.c`. + +### Entry points + +```c +int ag_backward(ag_tensor* loss); // scalar seed = 1.0 +int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient); +``` + +`ag_backward` requires `loss` to be a scalar (ndim == 0) and seeds with `ag_full_like(loss->value, 1.0f)`. `ag_backward_with_grad` allows non-scalar outputs with a user-supplied upstream gradient. + +### Algorithm step by step + +```mermaid +flowchart TD + A[ag_backward(loss)] --> B[Seed: create tensor of 1.0s matching loss shape] + B --> C[ag_backward_with_grad(loss, seed)] + + C --> D[collect_graph: DFS from loss through creator links] + D --> E[Assign graph_index to each tensor and append to tensors list] + D --> F[Append nodes in DFS post-order to nodes list] + + E --> G{graph_versions_match?} + F --> G + + G -->|No| H[Return 1 — stale graph rejected] + G -->|Yes| I[Allocate pass_gradients array] + + I --> J[Seed: pass_gradients[loss.graph_index] = clone of output_gradient] + + J --> K[Loop: for node_index = count-1 down to 0] + K --> L[node.backward(node, upstream_gradient, contributions)] + L --> M{For each input with requires_grad} + M -->|Yes| N[accumulate_pass_gradient: reduce to shape + add] + M -->|No| O[Free unused contribution] + + N --> P{More nodes?} + P -->|Yes| K + P -->|No| Q[merge_persistent_gradients] + + Q --> R[For each tensor with grad: merge = existing + pass gradient] + R --> S[Store merged gradients on ag_tensor.grad] + S --> T[Cleanup: reset graph_index, free lists] + T --> U[Return 0 — success] +``` + +### Graph collection: `collect_graph` + +The DFS in `collect_graph` (`autograd_backward.c:43-52`) walks from the loss tensor through `creator` links, depth-first: + +```c +static int collect_graph(ag_tensor* value, tensor_list* tensors, node_list* nodes) { + if (value == NULL || value->graph_index >= 0) return value == NULL; + value->graph_index = tensors->count; + if (append_tensor(tensors, value) != 0) return 1; + if (value->creator == NULL) return 0; // leaf + for (int i = 0; i < value->creator->input_count; ++i) { + if (collect_graph(value->creator->inputs[i], tensors, nodes) != 0) return 1; + } + return append_node(nodes, value->creator); +} +``` + +Each tensor gets a `graph_index` (used as an index into `pass_gradients`). Nodes are appended in post-order so that iterating `nodes` from back to front gives reverse topological order. + +### The backward loop + +From `autograd_backward.c:170-198`: + +```c +for (int node_index = nodes.count - 1; node_index >= 0; --node_index) { + ag_node* node = nodes.values[node_index]; + int gradient_index = node->output->graph_index; + + tensor* contributions[2] = {NULL, NULL}; + node->backward(node, pass_gradients[gradient_index], contributions); + + for (int input_index = 0; input_index < node->input_count; ++input_index) { + ag_tensor* input = node->inputs[input_index]; + if (!input->requires_grad) { t_free(contributions[input_index]); continue; } + accumulate_pass_gradient(&pass_gradients[input->graph_index], + contributions[input_index], input->value); + } +} +``` + +### Gradient accumulation + +`accumulate_pass_gradient` (`autograd_backward.c:103-118`) reduces a contribution to the target shape via `reduce_to_shape`, then adds it to the existing pass gradient for that tensor. If the destination is NULL (first contribution), it stores the reduced contribution directly. + +### Merging with persistent gradients + +After the backward loop, `merge_persistent_gradients` (`autograd_backward.c:120-145`) combines pass-local gradients with any previously accumulated `.grad` on each `ag_tensor`: + +```c +merged[i] = value->grad == NULL + ? t_clone(pass_gradients[i]) + : t_add(value->grad, pass_gradients[i]); +``` + +This is what makes repeated `ag_backward` calls accumulate: existing `.grad` values are added to new ones. + +### Comparison to PyTorch's `evaluate_function` + +PyTorch uses a task queue with worker threads. Each completed node enqueues its dependents. TensorLib uses a simpler single-threaded reverse iteration over the sorted node list. This is correct because the graph is a DAG and the post-order DFS guarantees dependencies are processed first. + +PyTorch's `Node::gradient_edge` design returns a single gradient per output; TensorLib's `ag_backward_fn` returns one gradient per input directly, avoiding the need for edge-indexed gradient lookups. + +--- + +## 7. Broadcast Gradient Reduction + +This is one of the most critical algorithms in the engine. When an input is broadcast during forward (e.g., adding a `[3]` vector to a `[2, 3]` matrix), the upstream gradient has the output shape, not the input shape. The gradient must be **reduced** back to the input's shape by summing along the broadcast dimensions. + +### The algorithm: `reduce_to_shape` + +From `autograd_backward.c:74-101`: + +```c +static tensor* reduce_to_shape(tensor* contribution, const tensor* target) { + tensor* current = contribution; + + // Step 1: Reduce extra leading dimensions + while (current->ndim > target->ndim) { + tensor* reduced = t_sum(current, 0); + t_free(current); + current = reduced; + } + + // Step 2: For each axis, if target has size 1 but current has size > 1, + // sum along that axis (keepdim) + for (int axis = 0; axis < target->ndim; ++axis) { + if (current->dims[axis] == target->dims[axis]) continue; + if (target->dims[axis] != 1 || current->dims[axis] == 1) { + t_free(current); return NULL; // shape mismatch + } + tensor* reduced = t_sum_keepdim(current, axis); + t_free(current); + current = reduced; + } + return current; +} +``` + +### Visual example + +Consider `a` with shape `[2, 3]` and `b` with shape `[3]`, computing `ag_add(a, b)`: + +```mermaid +flowchart LR + subgraph Forward + A["a [2,3]"] --> Add["ag_add"] + B["b [3]"] --> Add + Add --> Out["output [2,3]"] + end + + subgraph Backward + GOut["grad_out [2,3]"] --> ReduceA["Identity (no reduction)"] + GOut --> ReduceB["Sum along axis 0 → [3]"] + ReduceA --> GA["grad_a [2,3]"] + ReduceB --> GB["grad_b [3]"] + end +``` + +For `b`, the gradient `[2,3]` is reduced along axis 0 (where `b` was broadcast from `[3]` to `[2,3]`), summing the two rows to produce `[3]`. + +This reduction happens in `accumulate_pass_gradient`, which calls `reduce_to_shape` before accumulating the contribution into the pass gradient. + +--- + +## 8. Error Handling & Graph Consistency + +### Stale graph detection + +TensorLib uses **storage version counters** to detect when tensors have been modified between forward and backward. Every `Storage` has a monotonically increasing `version` field. When `tensor_mark_modified` is called, the version increments. + +During forward, `ag_make_result` captures the version of every input's storage and the output's storage: + +```c +node->input_versions[i] = inputs[i]->value->storage->version; +node->output_version = output->storage->version; +``` + +Before backward executes, `graph_versions_match` (`autograd_backward.c:54-72`) verifies that all captured versions still match: + +```c +static int graph_versions_match(const node_list* nodes) { + for (int node_index = 0; node_index < nodes->count; ++node_index) { + const ag_node* node = nodes.values[node_index]; + if (node->output->value->storage->version != node->output_version) + return 0; + for (int input_index = 0; input_index < node->input_count; ++input_index) { + if (node->inputs[input_index]->value->storage->version != + node->input_versions[input_index]) + return 0; + } + } + return 1; +} +``` + +If any version mismatches, `ag_backward_with_grad` returns `1` (failure) and **no gradients are modified**. + +### Transactional error handling + +The backward pass is transactional. `pass_gradients` is a local array; all computation happens there. Only after the entire backward loop succeeds does `merge_persistent_gradients` write to the actual `.grad` fields. If any step fails: + +```c +cleanup: + for (int i = 0; i < tensors.count; ++i) t_free(pass_gradients[i]); + free(pass_gradients); + for (int i = 0; i < tensors.count; ++i) tensors.values[i]->graph_index = -1; + // .grad fields are never touched — existing gradients are preserved + return status; // 1 = failure +``` + +This means a failed backward leaves all `.grad` fields exactly as they were, even if previous backward calls had accumulated gradients. + +### When does staleness happen? + +Staleness occurs when: + +1. A tensor's underlying `Storage` data is modified (via direct write + `tensor_mark_modified`). +2. A view's storage is modified through an alias. +3. The `ag_detach` alias is mutated. + +The engine detects this and refuses to compute gradients on a potentially corrupted graph. After a stale rejection, you can call `ag_zero_grad_all` and rebuild the graph if needed. + +--- + +## 9. API Reference + +### Tensor Construction & Lifetime + +| Function | Signature | Description | +|----------|-----------|-------------| +| `ag_from_owned_tensor` | `ag_tensor* ag_from_owned_tensor(tensor* value, int requires_grad)` | Wraps a raw tensor, taking ownership. Returns NULL on invalid metadata (frees value). | +| `ag_detach` | `ag_tensor* ag_detach(const ag_tensor* value)` | Creates a zero-copy leaf alias with `requires_grad=0`, sharing storage. | +| `ag_tensor_retain` | `void ag_tensor_retain(ag_tensor* value)` | Increments reference count. | +| `ag_tensor_release` | `void ag_tensor_release(ag_tensor* value)` | Decrements reference count; frees when zero. | +| `ag_node_retain` | `void ag_node_retain(ag_node* node)` | Increments node reference count. | +| `ag_node_release` | `void ag_node_release(ag_node* node)` | Decrements node reference count; frees when zero. | + +### Binary Operations + +| Function | Signature | +|----------|-----------| +| `ag_add` | `ag_tensor* ag_add(const ag_tensor* a, const ag_tensor* b)` | +| `ag_sub` | `ag_tensor* ag_sub(const ag_tensor* a, const ag_tensor* b)` | +| `ag_mul` | `ag_tensor* ag_mul(const ag_tensor* a, const ag_tensor* b)` | +| `ag_div` | `ag_tensor* ag_div(const ag_tensor* a, const ag_tensor* b)` | +| `ag_mul_scalar` | `ag_tensor* ag_mul_scalar(const ag_tensor* value, float scalar)` | +| `ag_div_scalar` | `ag_tensor* ag_div_scalar(const ag_tensor* value, float scalar)` | + +### Unary Operations + +| Function | Signature | +|----------|-----------| +| `ag_neg` | `ag_tensor* ag_neg(const ag_tensor* value)` | +| `ag_exp` | `ag_tensor* ag_exp(const ag_tensor* value)` | +| `ag_log` | `ag_tensor* ag_log(const ag_tensor* value)` | +| `ag_pow` | `ag_tensor* ag_pow(const ag_tensor* value, float exponent)` | +| `ag_sqrt` | `ag_tensor* ag_sqrt(const ag_tensor* value)` | +| `ag_relu` | `ag_tensor* ag_relu(const ag_tensor* value)` | +| `ag_sigmoid` | `ag_tensor* ag_sigmoid(const ag_tensor* value)` | +| `ag_tanh` | `ag_tensor* ag_tanh(const ag_tensor* value)` | +| `ag_gelu` | `ag_tensor* ag_gelu(const ag_tensor* value)` | + +### Reductions + +| Function | Signature | +|----------|-----------| +| `ag_sum` | `ag_tensor* ag_sum(const ag_tensor* value, int dim, int keepdim)` | +| `ag_mean` | `ag_tensor* ag_mean(const ag_tensor* value, int dim, int keepdim)` | +| `ag_max` | `ag_tensor* ag_max(const ag_tensor* value, int dim, int keepdim)` | + +### Views + +| Function | Signature | +|----------|-----------| +| `ag_reshape` | `ag_tensor* ag_reshape(const ag_tensor* value, int new_ndim, const int* new_dims)` | +| `ag_transpose` | `ag_tensor* ag_transpose(const ag_tensor* value, int dim0, int dim1)` | +| `ag_slice` | `ag_tensor* ag_slice(const ag_tensor* value, int dim, int start, int end)` | +| `ag_expand` | `ag_tensor* ag_expand(const ag_tensor* value, int new_ndim, const int* new_dims)` | + +### Matmul & Gather + +| Function | Signature | +|----------|-----------| +| `ag_matmul` | `ag_tensor* ag_matmul(const ag_tensor* a, const ag_tensor* b)` | +| `ag_gather_rows` | `ag_tensor* ag_gather_rows(const ag_tensor* table, const tensor* indices)` | + +Note: `ag_gather_rows` takes a raw `tensor*` for indices, not an `ag_tensor*`. Indices are not differentiable. + +### Backward & Gradient Management + +| Function | Signature | Description | +|----------|-----------|-------------| +| `ag_backward` | `int ag_backward(ag_tensor* loss)` | Backward from a scalar loss, seeding with 1.0. Returns 0 on success, 1 on failure. | +| `ag_backward_with_grad` | `int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient)` | Backward with explicit upstream gradient. Output must have `requires_grad=1`. | +| `ag_zero_grad` | `void ag_zero_grad(ag_tensor* value)` | Frees and NULLs a single tensor's `.grad`. | +| `ag_zero_grad_all` | `void ag_zero_grad_all(ag_tensor* root)` | Zeros `.grad` on every tensor reachable from root through the graph. | + +All forward operations return `NULL` on invalid arguments (NULL inputs, incompatible shapes, out-of-bounds dimensions, etc.). + +--- + +## 10. Test Coverage + +The autograd test suite comprises 9 test files in `tests/unit/autograd/`: + +| Test File | Focus | Key Patterns | +|-----------|-------|-------------| +| `test_autograd_core.c` | `ag_tensor` construction, `ag_detach`, retain/release lifecycle | Verifies ref counting, detached aliases sharing storage, null safety | +| `test_autograd_ops.c` | Binary ops (add/sub/mul/div), unary ops (neg/exp/log/pow/sqrt/relu/sigmoid/tanh/gelu), scalar ops | Local backward gradient checks, finite-difference verification, broadcast forward/backward, IEEE edge cases (NaN, Inf) | +| `test_autograd_view.c` | Reshape, transpose, slice, expand | Local backward shape verification, gradient scatter/gather correctness | +| `test_autograd_reduc.c` | Sum, mean, max reductions | Expand-backward for sum/mean, tie-splitting for max, NaN propagation | +| `test_autograd_matmul.c` | Matrix-matrix, vector-vector, vector-matrix, matrix-vector, batched matmul | Local backward checks, finite-difference validation, view-based matmul | +| `test_autograd_gather.c` | `ag_gather_rows` forward and backward | Scatter-add with duplicate indices, invalid index rejection | +| `test_autograd_backward.c` | Full backward pass through composed graphs | Chain gradients, shared DAG accumulation, unbroadcast verification, seeded backward, repeated accumulation, stale graph rejection, zero_grad | +| `test_autograd_integration.c` | End-to-end gradient correctness | Central-difference finite-difference verification for composed graphs (mul→exp→log→sum→mean), view chains, expand chains, lifecycle stress tests (500 iterations, MSVC debug heap leak detection) | +| `test_autograd_public_contract.c` | Comprehensive public API contract | Node ownership, graph omission for untracked tensors, broadcasting, local backward ownership, view gradients, reductions, matmul, accumulation, seed validation, stale graph rejection, null/invalid argument handling | + +### Key testing patterns + +**Finite-difference gradient checking**: The integration and ops tests compute numerical gradients via central differences `(f(x+ε) - f(x-ε)) / 2ε` and compare against autograd's analytical gradients, typically within `1e-3` to `1e-5` tolerance. + +**Transactionality testing**: Tests verify that modifying a tensor after forward causes `ag_backward` to return 1 and leaves existing gradients unchanged. + +**Memory leak detection**: On MSVC debug builds, `_CrtMemCheckpoint` / `_CrtMemDifference` verify zero leaked bytes after 500-iteration stress tests. + +--- + +## 11. Design Decisions & Tradeoffs + +### Dynamic graph (like PyTorch) vs static graph (like TensorFlow 1.x) + +TensorLib uses a **dynamic (define-by-run) graph**, matching PyTorch's approach. This means: + +- **Pros**: Natural control flow, simpler debugging (standard C debugger works), no session/graph compilation overhead, incremental memory management. +- **Cons**: No graph-level optimization (common subexpression elimination, operator fusion), no automatic batching across iterations. + +A static graph approach (like TensorFlow 1.x or ggml's `ggml_cgraph`) would enable graph optimizations but would require separating graph construction from execution, making the C API significantly more complex. + +### Manual reference counting vs garbage collection + +TensorLib uses **manual reference counting** with explicit `retain`/`release` calls. This is the natural choice for a C99 library: + +- **Pros**: Deterministic deallocation, no GC pauses, simple implementation, portable across all C99 compilers. +- **Cons**: Users must carefully release tensors; cycles would leak (though the DAG structure of autograd graphs prevents reference cycles in practice). + +The reference count on `ag_tensor` ensures that tensors shared between multiple consumers (e.g., an input used by two different operations) are kept alive until all consumers release them. + +### Ownership and reference count flow + +When a tensor is used as input to two operations, both nodes retain it: + +```mermaid +flowchart TD + A["ag_tensor a
ref_count=1"] -->|"ag_make_result retains"| N1["ag_node: add(a, c)
inputs[0]=a"] + A -->|"ag_make_result retains"| N2["ag_node: mul(a, b)
inputs[0]=a"] + N1 -->|"creator"| O1["ag_tensor: output1
ref_count=1"] + N2 -->|"creator"| O2["ag_tensor: output2
ref_count=1"] + + style A fill:#f9f,stroke:#333,stroke-width:2px + style O1 fill:#9f9,stroke:#333,stroke-width:2px + style O2 fill:#9f9,stroke:#333,stroke-width:2px +``` + +Releasing `output1` decrements `a`'s refcount from 3 to 2 (the `add` node releases its retained `a`). Releasing `output2` decrements it to 1. Only when the user finally releases the original `a` does its refcount reach 0 and its storage is freed. + +### Storage version counters vs PyTorch's version counter + +Both TensorLib and PyTorch use monotonically increasing version counters on storage to detect mutations. The mechanism is essentially identical: + +- TensorLib: `Storage.version` incremented by `tensor_mark_modified`. +- PyTorch: Variable version counter incremented on in-place modification. + +TensorLib checks versions for **both inputs and outputs** at every node, while PyTorch primarily checks input versions. This provides comprehensive staleness detection including cases where an intermediate output is modified. + +### Why 23 operations specifically? + +The 23 operations represent the **minimal set** needed for common deep learning workloads: + +- **4 binary arithmetic** (add, sub, mul, div) + 2 scalar variants +- **8 unary** (neg, exp, log, pow, sqrt, relu, sigmoid, tanh, gelu) — covers all standard activation functions +- **1 matmul** — the core linear algebra primitive +- **3 reductions** (sum, mean, max) — covers loss computation and normalization +- **4 views** (reshape, transpose, slice, expand) — covers tensor manipulation +- **1 gather** (gather_rows) — covers embedding lookups + +This set can express fully-connected layers, convolution (via im2col + matmul), attention mechanisms, layer normalization, and all standard loss functions. Operations like softmax can be composed from exp, sum, and div. See [neural_network_modules.md](./neural_network_modules.md) for how the NN layer uses autograd to build training loops. + +### Why `ag_make_result` as the single graph-construction point + +Every forward operation funnels through `ag_make_result`, which centralizes: + +1. Checking if any input requires gradients. +2. Allocating the node and result tensor. +3. Recording storage versions. +4. Retaining inputs. +5. Handling allocation failures transactionally. + +This eliminates duplicated graph-construction logic across 23 operations and ensures consistent error handling. + +--- + +## Source File Layout + +| File | Purpose | +|------|---------| +| `include/tensorlib/autograd.h` | Public API: `ag_tensor`, `ag_node`, all `ag_*` functions | +| `include/tensorlib/autograd_internal.h` | Internal helpers: `ag_make_result`, `ag_full_like` | +| `src/autograd/autograd_core.c` | Construction, lifetime, detach, `ag_make_result` | +| `src/autograd/autograd_ops.c` | Binary ops (add/sub/mul/div), scalar ops, unary ops (neg/exp/log/pow/sqrt/relu/sigmoid/tanh/gelu) | +| `src/autograd/autograd_view.c` | View ops (reshape, transpose, slice, expand) | +| `src/autograd/autograd_reduc.c` | Reduction ops (sum, mean, max) | +| `src/autograd/autograd_matmul.c` | Matrix multiplication with vector/batch support | +| `src/autograd/autograd_gather.c` | `gather_rows` with scatter-add backward | +| `src/autograd/autograd_backward.c` | Backward pass: graph collection, version check, traversal, accumulation | diff --git a/docs/decoder_implementation.md b/docs/decoder_implementation.md new file mode 100644 index 0000000..f2826d4 --- /dev/null +++ b/docs/decoder_implementation.md @@ -0,0 +1,990 @@ +# Decoder Implementation + +TensorLib's decoder is a complete, from-scratch implementation of a GPT-style +causal transformer decoder in C99. It composes token embeddings, positional +embeddings, configurable-depth pre-norm residual blocks with multi-head causal +self-attention, a final layer normalization, and a language-model projection +head into a single module that maps token IDs to next-token logits. + +The decoder is the highest-level building block in TensorLib's neural network +stack. It depends on the autograd engine for automatic differentiation, the +tensor library for storage and shape manipulation, and the individual neural +network primitives (linear, embedding, layer norm, attention, dropout) for its +internal layers. See [autograd_engine.md](./autograd_engine.md) for how +gradients flow through the computation graph, and +[tensor_mechanics.md](./tensor_mechanics.md) for tensor allocation and layout +details. + +## 1. Overview + +TensorLib's `nn_decoder` implements a GPT-2-style transformer decoder suitable +for autoregressive language modeling. The design is intentionally close to +[karpathy/nanoGPT](https://github.com/karpathy/nanoGPT) and +[PyTorch's TransformerDecoder](https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/transformer.py), +but written entirely in portable C99 with no external dependencies beyond the +standard library. + +Key characteristics: + +- **Configurable depth, width, and head count** via a single `nn_decoder_config` + struct. +- **Pre-norm residual blocks** (LayerNorm before attention and MLP) matching + GPT-2's design for stable deep training. +- **Fused QKV projection** in multi-head attention for better cache locality. +- **Learned positional embeddings** (not sinusoidal or rotary). +- **Causal masking** via additive lower-triangular mask with `-INFINITY` before + softmax. +- **Automatic differentiation** through the entire forward pass, enabling + `ag_backward` to populate gradients for every trainable parameter. +- **Checkpoint save/load** for model parameters, optimizer state, and RNG state. + +The decoder is used as the sole model in the +[TinyLM example](../examples/tiny_lm/README.md), which trains a ~1.9M +parameter byte-level language model on Shakespeare text. + +## 2. Architecture Overview + +The full decoder stack processes integer token IDs shaped `[batch, sequence]` +and produces logits shaped `[batch, sequence, vocabulary_size]`: + +```mermaid +flowchart TD + A["Input token IDs
[B, T]"] --> B["Token Embedding
[B, T, C]"] + A --> C["Position IDs
[0, 1, ..., T-1]"] + C --> D["Positional Embedding
[T, C]"] + B --> E["Add token + positional
[B, T, C]"] + D --> E + E --> F["Decoder Block 0"] + F --> G["Decoder Block 1"] + G --> H["..."] + H --> I["Decoder Block N-1"] + I --> J["Final LayerNorm
[B, T, C]"] + J --> K["LM Head (Linear)
[B, T, V]"] + K --> L["Output logits
[B, T, V]"] +``` + +Where: +- **B** = batch size +- **T** = sequence length (must be ≤ `context_length`) +- **C** = `channels` (model dimension, `d_model`) +- **V** = `vocabulary_size` +- **N** = `layer_count` (number of decoder blocks) +- **H** = `head_count` (number of attention heads) +- **D** = `channels / head_count` (per-head dimension) + +The decoder maps cleanly to GPT-2's architecture choices. GPT-2 Small uses +`n_layer=12, n_head=12, n_embd=768`; TensorLib's TinyLM uses +`layer_count=4, head_count=6, channels=192`. The only GPT-2 design elements +not present are weight-tying (embedding and LM head share weights) and a +separate key-value cache for efficient autoregressive inference. + +The `nn_decoder_config` struct encodes all hyperparameters: + +```c +struct nn_decoder_config { + int vocabulary_size; + int context_length; + int channels; + int head_count; + int layer_count; + float dropout_probability; + float layer_norm_epsilon; +}; +``` + +Validation in `decoder_config_valid` ensures `channels % head_count == 0`, +all counts are positive, dropout is in `[0, 1)`, and epsilon is finite and +positive. + +## 3. Causal Mask + +Autoregressive generation requires that each token position can only attend to +itself and earlier positions. TensorLib enforces this with an additive +lower-triangular mask applied to the raw attention scores before softmax. + +### Implementation + +The mask is constructed in `causal_mask.c` (`src/nn/causal_mask.c:5-36`): + +```c +ag_tensor* nn_apply_causal_mask(const ag_tensor* scores) +{ + int sequence = scores->value->dims[scores->value->ndim - 1]; + int dims[2] = { sequence, sequence }; + tensor* mask_value = t_alloc(2, dims); + for (int row = 0; row < sequence; ++row) { + for (int column = 0; column < sequence; ++column) { + mask_value->storage->data[row * sequence + column] = + column <= row ? 0.0f : -INFINITY; + } + } + ag_tensor* mask = ag_from_owned_tensor(mask_value, 0); + ag_tensor* result = ag_add(scores, mask); + ag_tensor_release(mask); + return result; +} +``` + +For a sequence length of 5, the mask matrix looks like: + +``` + 0.0 0.0 0.0 0.0 0.0 + 0.0 0.0 0.0 0.0 0.0 + 0.0 0.0 0.0 0.0 0.0 + 0.0 0.0 0.0 0.0 0.0 + 0.0 0.0 0.0 0.0 0.0 +``` + +Wait — that's the allowed positions. The *mask* that gets **added** is: + +``` + 0.0 0.0 0.0 0.0 0.0 + 0.0 0.0 0.0 0.0 0.0 + 0.0 0.0 0.0 0.0 0.0 + 0.0 0.0 0.0 0.0 0.0 + 0.0 0.0 0.0 0.0 0.0 +``` + +The **mask** itself has `0.0` on and below the diagonal, and `-INFINITY` above: + +``` + 0.0 -INF -INF -INF -INF + 0.0 0.0 -INF -INF -INF + 0.0 0.0 0.0 -INF -INF + 0.0 0.0 0.0 0.0 -INF + 0.0 0.0 0.0 0.0 0.0 +``` + +After `scores + mask`, every position above the diagonal becomes `-INFINITY`. +After softmax, these positions become `0.0` probability, preventing information +leakage from future tokens. + +```mermaid +flowchart LR + S["Raw scores
[B, H, T, T]"] --> M["nn_apply_causal_mask"] + M --> A["Masked scores
(-INF above diagonal)"] + A --> SM["softmax"] + SM --> P["Attention weights
(upper triangle = 0)"] +``` + +### Comparison to PyTorch + +PyTorch provides `torch.nn.Transformer.generate_square_subsequent_mask(sz)` +which returns the same lower-triangular structure but uses `-inf` as the +masked value. TensorLib's implementation is equivalent: it uses C's +`-INFINITY` from `` and applies the mask additively rather than +using `masked_fill_`. The effect is identical — `softmax(score + (-inf)) = 0`. + +## 4. Decoder Block (Pre-Norm Residual) + +Each decoder block implements the pre-norm residual architecture used by GPT-2 +and adopted by most modern transformers. The block contains two sub-layers +(attention and MLP), each preceded by a LayerNorm and wrapped in a skip +connection. + +### Forward Pass + +The forward pass in `decoder_block.c` (`src/nn/decoder_block.c:162-215`): + +```mermaid +flowchart TD + X["Input x"] --> LN1["LayerNorm (attention_norm)"] + LN1 --> MH["Multi-Head Attention"] + MH --> R1["Residual: x + MH(LN1(x))"] + X --> R1 + R1 --> LN2["LayerNorm (mlp_norm)"] + LN2 --> FC1["Linear (mlp_input): C → 4C"] + FC1 --> G["GELU Activation"] + G --> FC2["Linear (mlp_output): 4C → C"] + FC2 --> DO["Dropout"] + DO --> R2["Residual: r1 + Dropout(FC2(G(FC1(LN2(r1)))))"] + R1 --> R2 + R2 --> Y["Output"] +``` + +In code: + +```c +// Pre-norm attention sub-layer +normalized_attention = nn_layer_norm_forward(block->attention_norm, input); +attention_output = nn_multihead_attention_forward(block->attention, normalized_attention); +attention_residual = ag_add(input, attention_output); + +// Pre-norm MLP sub-layer +normalized_mlp = nn_layer_norm_forward(block->mlp_norm, attention_residual); +hidden = nn_linear_forward(block->mlp_input, normalized_mlp); // C → 4C +activated = ag_gelu(hidden); +projected = nn_linear_forward(block->mlp_output, activated); // 4C → C +dropped = nn_dropout_forward(block->mlp_dropout, projected); +result = ag_add(attention_residual, dropped); +``` + +### Why Pre-Norm? + +Pre-norm places the LayerNorm *before* the sub-layer computation rather than +after. This is the design choice made by GPT-2 (see the original +[Radford et al., 2019](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf)) +and differs from the original Transformer's post-norm. + +The key advantage is **training stability**: gradients flow through the +residual connection unimpeded by the normalization layer, which prevents +gradient magnitude from degrading in deep networks. In post-norm, the +LayerNorm sits directly in the gradient path, which can cause training +instability requiring careful learning rate warmup. Pre-norm generally +eliminates this requirement. + +### Comparison to PyTorch + +PyTorch's `TransformerEncoderLayer` defaults to post-norm in older versions +but offers `norm_first=True` (added in PyTorch 1.11) for pre-norm. TensorLib +uses pre-norm unconditionally, matching the GPT-2 convention. + +### Sub-components + +Each block registers 6 child modules: + +| Child | Type | Description | +|-------|------|-------------| +| `attention_norm` | `nn_layer_norm` | Pre-attention normalization | +| `attention` | `nn_multihead_attention` | Causal self-attention | +| `mlp_norm` | `nn_layer_norm` | Pre-MLP normalization | +| `mlp_input` | `nn_linear` | FFN expansion layer (`C → 4C`) | +| `mlp_output` | `nn_linear` | FFN contraction layer (`4C → C`) | +| `mlp_dropout` | `nn_dropout` | Dropout after FFN output | + +The hidden width is always `channels * 4`, matching the standard transformer +ratio from the original "Attention Is All You Need" paper and GPT-2. + +## 5. Multi-Head Attention (Deep Dive) + +Multi-head attention is the core mechanism that allows the decoder to attend to +all previous token positions in parallel. TensorLib implements causal +self-attention with a fused QKV projection for computational efficiency. + +### Fused QKV Projection + +Rather than maintaining three separate weight matrices `W_Q`, `W_K`, `W_V`, +TensorLib uses a single fused parameter of shape `[3, C, C]` and a fused bias +of shape `[3, C]`. This reduces the number of memory allocations and improves +cache locality during the projection. + +```mermaid +sequenceDiagram + participant I as Input [B,T,C] + participant F as Fused QKV
[3,C,C] × [1,B,T,C] + participant S as Split & Reshape + participant A as Scaled Dot-Product
Attention + participant O as Output Projection + + I->>F: ag_matmul(reshape(input), reshape(qkv_weight)) + F->>F: Add qkv_bias + F->>S: Slice [0]=Q, [1]=K, [2]=V + S->>S: Reshape to [B,T,H,D] and transpose + S->>A: scores = Q · K^T / sqrt(D) + A->>A: Apply causal mask (-INF above diagonal) + A->>A: softmax → attention weights + A->>A: context = weights · V + A->>O: Concatenate heads, reshape to [B,T,C] + O->>O: Linear output projection [C,C] + O->>O: Output dropout +``` + +### Forward Pass + +The attention forward pass in `multihead_attention.c` +(`src/nn/multihead_attention.c:227-309`): + +1. **Fused QKV projection**: Reshape input to `[1, B, T, C]`, matmul with + weight `[3, 1, C, C]`, add bias `[3, 1, 1, C]`, producing QKV of shape + `[3, B, T, C]`. + +2. **Split into heads**: Slice Q, K, V along the first dimension, reshape each + to `[B, T, H, D]`, then transpose to `[B, H, T, D]` for batched attention. + +3. **Scaled dot-product**: `scores = Q · K^T / sqrt(D)` produces + `[B, H, T, T]`. + +4. **Causal mask**: `nn_apply_causal_mask` adds the lower-triangular mask, + setting future positions to `-INFINITY`. + +5. **Softmax**: `nn_softmax` converts masked scores to attention weights. + +6. **Context**: `context = weights · V` produces `[B, H, T, D]`. + +7. **Merge heads**: Transpose back to `[B, T, H, D]`, reshape to `[B, T, C]`. + +8. **Output projection**: Linear `[C, C]` followed by dropout. + +The scaling factor `1/sqrt(D)` prevents the dot-product magnitude from growing +with head dimension, which would push softmax into regions with tiny gradients. + +### Comparison to FlashAttention and ggml + +TensorLib implements the standard attention algorithm. For comparison: + +| Implementation | Approach | Key Difference | +|---|---|---| +| TensorLib | Standard matmul + softmax | Reference implementation, full autograd support | +| [FlashAttention](https://github.com/Dao-AILab/flash-attention) (Dao et al.) | Tiling + online softmax | IO-aware, avoids materializing full `[B,H,T,T]` matrix | +| [ggml](https://github.com/ggerganov/llama.cpp) `ggml_flash_attn` | Fused kernel | CPU/GPU fused kernel, no autograd | +| [PyTorch](https://github.com/pytorch/pytorch/blob/main/torch/nn/functional.py) `scaled_dot_product_attention` | Dispatches to Flash/MEMORY backend | Automatic backend selection | + +TensorLib's attention is designed for **correctness and clarity** first. It +materializes the full score matrix, making autograd straightforward. For +production inference, FlashAttention or ggml kernels would be more memory- +and compute-efficient. + +### QKV Weight Initialization + +The fused QKV weight is initialized with uniform samples from +`[-sqrt(3/C), sqrt(3/C)]` (a form of LeCun/Xavier uniform), and the bias is +initialized to zero. This is done with a manual loop in +`nn_multihead_attention_create` after the parameter is created with +`NN_INIT_ZERO`: + +```c +scale = sqrtf(3.0f / (float)channels); +for (int index = 0; + index < tensor_numel(attention->qkv_weight->value->value); + ++index) { + attention->qkv_weight->value->value->storage->data[index] = + nn_rng_uniform(rng, -scale, scale); +} +``` + +## 6. MLP Block + +Each decoder block contains a two-layer feedforward network (MLP) following the +attention sub-layer. The MLP is an in-place transformation of each token +position independently. + +### Architecture + +``` +Input [B, T, C] + → Linear [C, 4C] (mlp_input) + → GELU activation + → Linear [4C, C] (mlp_output) + → Dropout +Output [B, T, C] +``` + +The hidden dimension `4C` (stored as `hidden_width = channels * 4` in +`decoder_block.c:54`) follows the standard 4x expansion ratio from the original +Transformer paper. This ratio appears in GPT-2, GPT-3, and virtually all +decoder-only transformers. + +### GELU Activation + +TensorLib uses the Gaussian Error Linear Unit (GELU) activation, computed via +`ag_gelu` in the autograd graph. GELU is the standard activation in GPT-2 and +most modern transformers, chosen for its smooth gradient properties compared +to ReLU. + +### Comparison to Transformer FFN Design + +The standard transformer FFN is defined as: + +``` +FFN(x) = W_out · activation(W_in · x + b_in) + b_out +``` + +TensorLib's implementation matches this exactly, with bias terms enabled in +both linear layers. The only omission relative to GPT-2 is that the output +projection in some GPT-2 implementations omits bias on the second linear +layer — TensorLib always includes it. + +See [neural_network_modules.md](./neural_network_modules.md) for details on +the `nn_linear` and activation function primitives. + +## 7. Token & Positional Embeddings + +### Token Embedding + +The token embedding maps integer token IDs (represented as floats in the +autograd tensor) to dense vectors of dimension `channels`. It is a standard +lookup table: + +```c +// embedding.c:86-94 +ag_tensor* nn_embedding_forward(const nn_embedding* layer, + const ag_tensor* indices) { + return ag_gather_rows(layer->weight->value, indices->value); +} +``` + +The weight matrix has shape `[vocabulary_size, channels]` and is initialized +with Xavier uniform. Token IDs must be non-negative integers less than +`vocabulary_size`, stored as float values in the input tensor. + +The embedding layer does not require gradients on the input (token IDs are +indices, not differentiable values), but the weight matrix is a trainable +parameter. + +### Positional Embedding + +Position information is added via a learned positional embedding table. The +positional embedding stores a `[context_length, channels]` weight matrix +(essentially an `nn_embedding` with `vocabulary_size = context_length`). + +During forward (`positional_embedding.c:84-118`): + +1. Generate position IDs `[0, 1, 2, ..., T-1]` as a float tensor. +2. Look up the corresponding rows from the positional table: `[T, C]`. +3. **Add** the positional vectors to the token embeddings element-wise. + +```c +// Generate position IDs +position_values = t_alloc(1, position_dims); +for (int position = 0; position < sequence; ++position) { + position_values->storage->data[position] = (float)position; +} +position_ids = ag_from_owned_tensor(position_values, 0); +positions = nn_embedding_forward(layer->table, position_ids); +result = ag_add(token_embeddings, positions); +``` + +### Combination + +Token and positional embeddings are combined by **addition** (not +concatenation), producing `[B, T, C]`. This matches the original Transformer +and GPT-2 convention. The learned positional approach is identical to GPT-2 +— sinusoidal positional encodings from the original paper are not used. + +## 8. Full Decoder Assembly + +### The `nn_decoder` Struct + +```c +struct nn_decoder { + nn_module base; + + nn_embedding* token_embedding; + nn_positional_embedding* positional_embedding; + nn_decoder_block** blocks; + size_t block_count; + nn_layer_norm* final_norm; + nn_linear* language_model_head; + + nn_decoder_config config; +}; +``` + +### Module Tree + +The decoder registers all sub-modules as children of its `nn_module` base, +forming a tree that enables recursive operations like `nn_module_set_training`, +`nn_module_zero_grad`, and parameter counting. + +```mermaid +flowchart TD + D["nn_decoder
(base: nn_module)"] --> TE["token_embedding
(nn_embedding)"] + D --> PE["positional_embedding
(nn_positional_embedding)"] + PE --> PET["table
(nn_embedding)"] + D --> B0["blocks[0]
(nn_decoder_block)"] + D --> B1["blocks[1]
(nn_decoder_block)"] + D --> BN["..."] + D --> FN["final_norm
(nn_layer_norm)"] + D --> LM["language_model_head
(nn_linear)"] + + B0 --> AN0["attention_norm
(nn_layer_norm)"] + B0 --> AT0["attention
(nn_multihead_attention)"] + B0 --> MN0["mlp_norm
(nn_layer_norm)"] + B0 --> MI0["mlp_input
(nn_linear)"] + B0 --> MO0["mlp_output
(nn_linear)"] + B0 --> MD0["mlp_dropout
(nn_dropout)"] + + AT0 --> QKV["qkv_weight [3,C,C]
(nn_parameter)"] + AT0 --> QKVB["qkv_bias [3,C]
(nn_parameter)"] + AT0 --> AO["output
(nn_linear)"] + AT0 --> AD["output_dropout
(nn_dropout)"] +``` + +### Forward Pass + +`nn_decoder_forward` (`src/nn/decoder.c:213-256`) orchestrates the full +pipeline: + +```c +ag_tensor* nn_decoder_forward(const nn_decoder* decoder, + const ag_tensor* token_ids) { + tokens = nn_embedding_forward(decoder->token_embedding, token_ids); + current = nn_positional_embedding_forward( + decoder->positional_embedding, tokens); + for (size_t index = 0; index < decoder->block_count; ++index) { + next = nn_decoder_block_forward(decoder->blocks[index], current); + ag_tensor_release(current); + current = next; + } + normalized = nn_layer_norm_forward(decoder->final_norm, current); + result = nn_linear_forward(decoder->language_model_head, normalized); + return result; +} +``` + +The final LayerNorm + linear head is equivalent to GPT-2's approach. GPT-2 +optionally ties the LM head weights with the token embedding weights (weight +tying); TensorLib does **not** tie weights, treating the LM head as an +independent linear projection. + +### Loss Computation + +`nn_decoder_loss` (`src/nn/decoder.c:258-295`) combines forward pass with +cross-entropy loss for convenient training: + +```c +ag_tensor* nn_decoder_loss(const nn_decoder* decoder, + const ag_tensor* token_ids, + const tensor* targets) { + logits = nn_decoder_forward(decoder, token_ids); + flattened_logits = ag_reshape(logits, 2, logit_dims); + flattened_targets = t_reshape(targets, 1, target_dims); + loss = nn_cross_entropy(flattened_logits, flattened_targets); + return loss; +} +``` + +Both `token_ids` and `targets` have shape `[B, T]`. The targets are +shifted-by-one: for input `[t0, t1, t2, ...]`, targets are +`[t1, t2, t3, ...]`. The loss function handles flattening to `[B*T, V]` +and `[B*T]` internally, computes log-softmax, selects the log-probability of +the correct class, negates, and averages. + +### Parameter Count Formula + +For a decoder with vocabulary size `V`, channels `C`, head count `H`, and +`N` layers: + +| Component | Parameters | +|-----------|-----------| +| Token embedding | `V × C` | +| Positional embedding | `L_ctx × C` (where `L_ctx` = context length) | +| Per decoder block: | | +| — attention_norm | `2C` (weight + bias) | +| — QKV weight | `3 × C × C` | +| — QKV bias | `3 × C` | +| — attention output | `C × C + C` | +| — mlp_norm | `2C` | +| — mlp_input | `C × 4C + 4C` | +| — mlp_output | `4C × C + C` | +| Block subtotal | `N × (18C² + 12C)` | +| Final norm | `2C` | +| LM head | `C × V + V` | +| **Total** | `V × C + L_ctx × C + N × (18C² + 12C) + 2C + C × V + V` | + +For the test configuration (`V=5, C=8, H=2, N=2, L_ctx=4`): +- Token embedding: 5 × 8 = 40 +- Positional embedding: 4 × 8 = 32 +- 2 blocks × (18×64 + 12×8) = 2 × 1248 = 2496 → wait, this includes + per-head splitting that doesn't add parameters. The actual count is + verified by the unit test as **30 parameter tensors** (not elements — + parameter count is the number of `nn_parameter` objects, not total floats). + +## 9. Training the Decoder + +### Loss Function + +The decoder uses cross-entropy loss for next-token prediction. The loss is +computed over all positions simultaneously — for each position `t`, the model +predicts the token at position `t+1`. The cross-entropy implementation +(`src/losses/classification.c:84-141`) computes: + +1. Numerically stable log-softmax over the vocabulary dimension. +2. Selection of the log-probability corresponding to the target class. +3. Negation to get per-position loss. +4. Mean reduction to scalar. + +### Optimizer: AdamW + +Training uses AdamW with gradient clipping, implemented in +`src/optim/adamw.c`. Key configuration: + +| Hyperparameter | Typical Value | Description | +|---|---|---| +| `learning_rate` | `3e-4` | Step size | +| `beta1` | `0.9` | First moment decay | +| `beta2` | `0.999` | Second moment decay | +| `epsilon` | `1e-8` | Numerical stability | +| `weight_decay` | `0.01` | L2 regularization | +| `max_grad_norm` | `1.0` | Gradient clipping threshold | + +The AdamW step (`src/optim/adamw.c:211-275`) applies bias correction and +decoupled weight decay. Gradient clipping computes the global L2 norm across +all parameters and rescales gradients if the norm exceeds `max_grad_norm`. + +### Training Loop + +```mermaid +sequenceDiagram + participant C as Training Loop + participant O as AdamW Optimizer + participant D as nn_decoder + participant A as Autograd + + C->>O: nn_adamw_zero_grad(optimizer) + C->>D: nn_decoder_loss(model, inputs, targets) + Note over D: Forward pass: embeddings → blocks → logits → cross-entropy + D-->>C: scalar loss + C->>A: ag_backward(loss) + Note over A: Backprop through entire graph
All parameters get .grad + C->>O: nn_adamw_step(optimizer) + Note over O: Clip gradients, update moments,
apply weight decay, update weights +``` + +From the TinyLM example (`examples/tiny_lm/tiny_lm.c:530-589`): + +```c +for (int step = 1; step <= options.steps; ++step) { + nn_adamw_zero_grad(optimizer); + make_batch(&corpus, 1, options.batch_size, step, &rng, &inputs, &targets); + loss = nn_decoder_loss(model, inputs, targets); + ag_backward(loss); + nn_adamw_step(optimizer); + ag_tensor_release(loss); +} +``` + +See [autograd_engine.md](./autograd_engine.md) for details on how +`ag_backward` traverses the computation graph and accumulates gradients. + +### Checkpointing + +The training loop periodically saves checkpoints containing model parameters, +AdamW first/second moments, step counts, and RNG state: + +```c +nn_checkpoint_save(options.checkpoint_path, &model->base, optimizer, &rng); +``` + +Checkpoints can be resumed with `--resume`, restoring all training state +transactionally — if validation fails, existing state is left intact. + +## 10. Inference / Text Generation + +### Autoregressive Generation + +The decoder generates text one token at a time. For each new token: + +1. Encode the current context as a `[1, T]` tensor of float token IDs. +2. Run `nn_decoder_forward` to get logits `[1, T, V]`. +3. Extract the logits for the **last position** only: `logits[0, T-1, :]`. +4. Sample from the distribution (top-k with temperature, or greedy). +5. Append the sampled token to the context buffer. +6. If the context exceeds `context_length`, shift the window left. + +The TinyLM example implements this in `generate()` (`examples/tiny_lm/tiny_lm.c:404-467`): + +```c +for (int generated = 0; generated < options.generate_count; ++generated) { + tensor* raw = t_alloc(2, dims); + for (size_t index = 0; index < context_size; ++index) { + raw->storage->data[index] = (float)context[index]; + } + input = ag_from_owned_tensor(raw, 0); + logits = nn_decoder_forward(model, input); + next = sample_next(logits, context_size, temperature, top_k, &rng); + fputc(next, stdout); + // Shift context window + if (context_size < TINY_LM_CONTEXT) { + context[context_size++] = (unsigned char)next; + } else { + memmove(context, context + 1, TINY_LM_CONTEXT - 1); + context[TINY_LM_CONTEXT - 1] = (unsigned char)next; + } +} +``` + +### Top-K Sampling + +The `sample_next` function performs top-k sampling: + +1. Sort all vocabulary logits by descending score. +2. Select the top `k` candidates. +3. Divide by temperature to sharpen or flatten the distribution. +4. Compute softmax weights for the top-k candidates. +5. Sample from the resulting categorical distribution. + +When `temperature = 0`, greedy decoding (argmax) is used. When `top_k = 0`, +the full vocabulary is considered. + +### KV-Caching + +TensorLib does **not** implement KV-caching. Each generation step recomputes +the full forward pass over the entire context window. This is correct but +computationally expensive for long sequences, since the cost per generated +token is `O(T²)` for attention. + +A KV-cache would store the projected keys and values from previous positions, +reducing per-step cost to `O(T)` for the attention layer. This is a common +optimization in production inference engines (e.g., vLLM, llama.cpp, TGI) but +adds significant implementation complexity to the attention module. + +## 11. Example: TinyLM + +The TinyLM example (`examples/tiny_lm/tiny_lm.c`) is a complete, end-to-end +byte-level language model trainer and text generator. + +### Architecture + +```c +enum { + TINY_LM_VOCABULARY = 256, // Byte-level vocabulary + TINY_LM_CONTEXT = 128, // Maximum sequence length + TINY_LM_CHANNELS = 192, // Model dimension + TINY_LM_HEADS = 6, // Attention heads + TINY_LM_LAYERS = 4 // Decoder blocks +}; +``` + +With these dimensions, TinyLM has approximately **1.9 million** trainable +parameters. The head width is `192 / 6 = 32` dimensions per head. + +### Training on Shakespeare + +The model trains on raw bytes (no tokenizer needed). The corpus is split +90/10 for training and validation. Key training settings: + +- Batch size: 1 (configurable) +- Learning rate: `3e-4` +- Weight decay: `0.01` +- Gradient clipping: `max_norm = 1.0` +- Dropout: `0.1` +- Optimizer: AdamW + +### Performance + +On an i5-13400F in Release mode: +- **~0.66 seconds** per batch-1 update +- **~193 byte tokens/second** +- 1,000 updates ≈ 11 minutes +- 50,000 updates ≈ 9 hours + +See [examples/tiny_lm/README.md](../examples/tiny_lm/README.md) for corpus +guidance, build instructions, and configuration options. + +## 12. API Reference + +### Decoder + +```c +nn_decoder* nn_decoder_create( + const char* name, + const nn_decoder_config* config, + nn_rng* rng +); + +void nn_decoder_destroy(nn_decoder* decoder); + +ag_tensor* nn_decoder_forward( + const nn_decoder* decoder, + const ag_tensor* token_ids +); + +ag_tensor* nn_decoder_loss( + const nn_decoder* decoder, + const ag_tensor* token_ids, + const tensor* targets +); +``` + +- `nn_decoder_create` — Allocates the decoder and all sub-modules. Returns + `NULL` on invalid config, OOM, or if `rng` is `NULL`. +- `nn_decoder_forward` — Maps `[B, T]` token IDs to `[B, T, V]` logits. + Token IDs must not require gradients. Sequence length must be ≤ + `context_length`. +- `nn_decoder_loss` — Forward pass followed by cross-entropy loss over all + positions. Returns a scalar `ag_tensor` suitable for `ag_backward`. +- `nn_decoder_destroy` — Recursively frees all sub-modules and parameters. + +### Decoder Block + +```c +nn_decoder_block* nn_decoder_block_create( + const char* name, + int channels, + int head_count, + float dropout_probability, + float layer_norm_epsilon, + nn_rng* rng +); + +void nn_decoder_block_destroy(nn_decoder_block* block); + +ag_tensor* nn_decoder_block_forward( + const nn_decoder_block* block, + const ag_tensor* input +); +``` + +Input must be `[B, T, C]`. Output has the same shape. + +### Multi-Head Attention + +```c +nn_multihead_attention* nn_multihead_attention_create( + const char* name, + int channels, + int head_count, + float dropout_probability, + nn_rng* rng +); + +void nn_multihead_attention_destroy(nn_multihead_attention* attention); + +ag_tensor* nn_multihead_attention_forward( + const nn_multihead_attention* attention, + const ag_tensor* input +); +``` + +Input must be `[B, T, C]` where `C` equals `channels`. Causal masking is +applied automatically. + +### Causal Mask + +```c +ag_tensor* nn_apply_causal_mask(const ag_tensor* scores); +``` + +Applies a lower-triangular additive mask to the last two dimensions of +`scores` (which must be square). Entries above the diagonal are set to +`-INFINITY`. + +### Module Utilities + +```c +size_t nn_module_parameter_count(const nn_module* module); +nn_parameter* nn_module_parameter_at(const nn_module* module, size_t index); +void nn_module_set_training(nn_module* module, int training); +int nn_module_is_training(const nn_module* module); +void nn_module_zero_grad(nn_module* module); +int nn_clip_grad_norm(nn_module* module, float max_norm, float* total_norm); +``` + +### Checkpointing + +```c +int nn_checkpoint_save(const char* path, const nn_module* module, + const nn_adamw* optimizer, const nn_rng* rng); +int nn_checkpoint_load(const char* path, nn_module* module, + nn_adamw* optimizer, nn_rng* rng); +``` + +## 13. Test Coverage + +### Decoder Tests (`tests/unit/nn/test_nn_decoder.c`) + +| Test | What It Verifies | +|------|-----------------| +| `test_topology_forward_loss_and_backward` | Block count, child count (6), parameter count (30 tensors), unique parameter names, forward output shape `[2,4,5]`, loss is scalar finite, backward populates all gradients | +| `test_validation_and_mode_propagation` | Rejects invalid configs (odd head count, 0 layers, dropout=1.0, NULL name/rng), rejects bad inputs (NaN, grad-tracked, too-long, wrong targets), training mode propagates to deepest sub-modules | +| `test_configurable_depth` | 3-layer decoder has `child_count=7`, `parameter_count=42`, correct block naming (`deep.blocks.2`), correct output shape | +| `test_tiny_batch_overfit` | Trains on a single sequence `[0,1,2,3]→[1,2,3,4]` for 300 steps, verifies loss < 0.05, verifies argmax predictions match targets | + +### Decoder Block Tests (`tests/unit/nn/test_nn_decoder_block.c`) + +| Test | What It Verifies | +|------|-----------------| +| `test_residual_forward_backward_and_topology` | 6 children, 12 parameter tensors, zero-residual identity test (output equals input when all weights are zero), gradient = 1.0 through skip connections, all parameters get gradients | +| `test_validation_training_and_lysis` | Rejects NULL name, non-divisible channels, zero epsilon, NULL rng, wrong input rank/width; eval mode freezes dropout RNG state | + +### Causal Mask Tests (`tests/unit/nn/test_nn_causal_mask.c`) + +| Test | What It Verifies | +|------|-----------------| +| `test_mask_and_softmax` | Above-diagonal entries become `-INFINITY`, on/below-diagonal unchanged; after softmax, above-diagonal is 0.0 and rows sum to 1.0; backward produces gradient = 1.0 everywhere | +| `test_invalid` | Rejects NULL, 1D vector, non-square 2D tensor | + +### Multi-Head Attention Tests (`tests/unit/nn/test_nn_multihead_attention.c`) + +| Test | What It Verifier | +|------|-----------------| +| `test_exact_forward_causality_and_backward` | With identity QKV and output projections, output matches reference attention implementation; changing future tokens does not affect earlier outputs (causality); backward gradient verified against numerical differentiation | +| `test_topology_validation_and_eval_rng` | Rejects non-divisible channels, 0 heads, NULL rng; rejects wrong input rank/width; eval mode freezes RNG state; 4 parameters, 2 children | + +## 14. Design Decisions & Tradeoffs + +### Learned vs Sinusoidal vs Rotary Positional Embeddings + +TensorLib uses **learned positional embeddings** (`nn_positional_embedding`), +matching GPT-2's design. This stores a trainable `[context_length, channels]` +matrix looked up by position index. + +**Alternatives considered:** + +- **Sinusoidal** (original Transformer): Fixed, no parameters, but requires + the model to learn the meaning of each frequency component. Learned embeddings + have been shown to work at least as well for fixed-length contexts. +- **RoPE** (Rotary Position Embeddings): Encodes relative position by rotating + query/key vectors. More parameter-efficient and generalizes to longer + sequences, but requires modifying the attention mechanism and prevents the + use of a simple embedding table. Would require significant changes to + `nn_multihead_attention`. + +The learned approach is the simplest to implement correctly in a from-scratch +C library and matches GPT-2's architecture exactly. + +### No KV-Cache + +TensorLib does not implement key-value caching for autoregressive inference. +Each forward pass recomputes all keys and values from scratch. + +**Why:** KV-cache adds complexity to the attention module (storing and +concatenating previous K/V tensors across calls) and interacts with the module +abstraction boundary. For a reference implementation focused on training +correctness, the simplicity of stateless attention is preferable. Production +inference engines (llama.cpp, vLLM) add KV-cache as an optimization layer on +top of the core attention logic. + +### Fused QKV vs Separate Projections + +The QKV projection uses a single `[3, C, C]` weight tensor rather than +separate `W_Q`, `W_K`, `W_V` matrices. + +**Why:** This reduces the number of parameter objects and memory allocations. +In practice, it also enables a single matmul operation to compute all three +projections simultaneously, improving cache locality. The output is sliced +along the leading dimension to obtain Q, K, V — an O(1) operation in terms +of data movement since TensorLib uses views for slicing. + +### Pre-Norm vs Post-Norm + +The decoder uses pre-norm (LayerNorm before each sub-layer) unconditionally. + +**Why:** Pre-norm is the standard in GPT-2 and all subsequent GPT models. +Empirically, pre-norm provides more stable training for deep networks because +gradients flow directly through the residual connections without being +modulated by the normalization layer. This eliminates the need for learning +rate warmup schedules that post-norm often requires. + +The tradeoff is that pre-norm may produce slightly lower final performance +than carefully tuned post-norm with warmup, but it is significantly more +robust to hyperparameter choices. + +### Memory Layout: Batch × Sequence × Hidden + +The decoder uses `[B, T, C]` layout throughout, with the hidden dimension +as the innermost (contiguous) dimension. + +**Why:** This layout ensures that the linear projection +`ag_matmul(input, weight^T)` operates on the last dimension, which is the +default behavior of the tensor library's matmul. It also means that +individual token vectors are contiguous in memory, which is favorable for +the embedding lookup and layer norm operations that normalize over the +hidden dimension. + +The attention module temporarily transposes to `[B, H, T, D]` for the +batched dot-product, then transposes back. This is a standard layout +choice shared by PyTorch, JAX, and most transformer implementations. + +### Weight Initialization + +The decoder uses: +- **Xavier uniform** for most linear layers and embeddings +- **Custom uniform `[-sqrt(3/C), sqrt(3/C)]`** for the QKV weight (functionally + equivalent to Xavier uniform for the QKV case) +- **Zero** for all biases +- **One** for LayerNorm weights, **zero** for LayerNorm biases + +This matches common practice and ensures that the initial output variance is +approximately preserved through each layer. diff --git a/docs/neural_network_modules.md b/docs/neural_network_modules.md new file mode 100644 index 0000000..db58caf --- /dev/null +++ b/docs/neural_network_modules.md @@ -0,0 +1,1010 @@ +# Neural Network Modules + +TensorLib's neural network component provides a complete, from-scratch deep learning module system in C99. It implements a composable module hierarchy with automatic parameter management, differentiable forward passes, loss functions, optimizers, and checkpoint serialization — all built on top of the autograd engine. + +> **See also:** [tensor_mechanics.md](./tensor_mechanics.md) for base tensor ops, [autograd_engine.md](./autograd_engine.md) for gradient computation, and [decoder_implementation.md](./decoder_implementation.md) for the full GPT-style decoder. + +--- + +## 1. Overview + +The `nn` component provides: + +- **Module system** — composable parent/child hierarchy with virtual dispatch via function pointers +- **Parameter system** — named, trainable tensors with automatic weight initialization +- **Layer library** — Linear, Embedding, PositionalEmbedding, LayerNorm, Dropout, MultiheadAttention, MLP, DecoderBlock, and a full causal language model Decoder +- **Loss functions** — numerically stable softmax, log-softmax, and cross-entropy +- **Optimizers** — SGD and AdamW (with decoupled weight decay and gradient clipping) +- **RNG** — deterministic splitmix64 PRNG for reproducible initialization and dropout +- **Checkpointing** — atomic binary save/load of model params, optimizer state, and RNG state + +### Comparison to PyTorch's `nn.Module` + +TensorLib's `nn_module` serves the same role as PyTorch's [`torch.nn.Module`](https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/module.py). The key differences: + +| Feature | PyTorch | TensorLib | +|---|---|---| +| Language | Python (with C++ extensions) | Pure C99 | +| Module composition | `self.linear = nn.Linear(...)` in `__init__` | `nn_module_register_child()` / `nn_module_register_parameter()` | +| Parameter access | `module.parameters()` (generator) | `nn_module_parameter_at()` (index-based) | +| Forward dispatch | `__call__` → `forward()` (Python override) | `nn_module_forward()` → function pointer callback | +| Train/eval mode | `module.train()` / `module.eval()` | `nn_module_set_training()` (recursive) | +| Memory management | Python garbage collector | Manual `create`/`destroy` pairs | +| Serialization | `torch.save` (pickle) | Custom binary format with magic `TLCKPT` | + +TensorLib also draws design inspiration from [ggml's `ggml_build_forward`](https://github.com/ggerganov/llama.cpp/blob/master/ggml.c) for its graph-based autograd approach, where each forward pass builds a computation graph that can later be traversed in reverse for backpropagation. + +--- + +## 2. Module System + +### The `nn_module` Struct + +Every layer in TensorLib embeds an `nn_module` as its first member, enabling C-style polymorphism through composition (the "first-member" pattern): + +```c +struct nn_module { + const char* type_name; // e.g. "Linear", "LayerNorm" + char* name; // caller-assigned instance name + + nn_module_forward_fn forward; // virtual dispatch callback + nn_module_destroy_fn destroy; // virtual destructor + + nn_parameter** parameters; // direct parameters (weight tensors) + size_t parameter_count; + size_t parameter_capacity; + + nn_module** children; // sub-modules + size_t child_count; + size_t child_capacity; + + int training; // 1 = train, 0 = eval +}; +``` + +Derived layer types embed `nn_module` as their first field, allowing safe casting: + +```c +struct nn_linear { + nn_module base; // must be first + nn_parameter* weight; + nn_parameter* bias; + int in_features; + int out_features; + int use_bias; +}; +``` + +This mirrors the inheritance pattern in C++ where `nn_linear*` can be treated as `nn_module*`. + +### Module Composition + +Modules form a tree. Parent modules own references to their children and parameters: + +```c +nn_linear* linear = nn_linear_create("proj", 64, 128, 1, ...); +nn_module_register_child(&parent->base, &linear->base); +``` + +Child registration includes cycle detection — `nn_module_register_child` rejects additions that would create a cycle (the target module already contains the parent). + +```mermaid +classDiagram + class nn_module { + +const char* type_name + +char* name + +nn_module_forward_fn forward + +nn_module_destroy_fn destroy + +nn_parameter** parameters + +nn_module** children + +int training + } + class nn_linear { + +nn_parameter* weight + +nn_parameter* bias + +int in_features + +int out_features + } + class nn_layer_norm { + +nn_parameter* weight + +nn_parameter* bias + +int normalized_width + +float epsilon + } + class nn_multihead_attention { + +nn_parameter* qkv_weight + +nn_parameter* qkv_bias + +nn_linear* output + +nn_dropout* output_dropout + } + class nn_decoder_block { + +nn_layer_norm* attention_norm + +nn_multihead_attention* attention + +nn_layer_norm* mlp_norm + +nn_linear* mlp_input + +nn_linear* mlp_output + +nn_dropout* mlp_dropout + } + class nn_decoder { + +nn_embedding* token_embedding + +nn_positional_embedding* positional_embedding + +nn_decoder_block** blocks + +nn_layer_norm* final_norm + +nn_linear* language_model_head + } + nn_module <|-- nn_linear + nn_module <|-- nn_layer_norm + nn_module <|-- nn_multihead_attention + nn_module <|-- nn_decoder_block + nn_module <|-- nn_decoder + nn_decoder_block *-- nn_multihead_attention + nn_decoder_block *-- nn_layer_norm + nn_decoder *-- nn_decoder_block + nn_decoder *-- nn_embedding +``` + +### Forward Dispatch + +`nn_module_forward` invokes the function pointer set during construction: + +```c +ag_tensor* nn_module_forward(const nn_module* module, const ag_tensor* input) +{ + if (module == NULL || input == NULL || module->forward == NULL) return NULL; + return module->forward(module, input); +} +``` + +Each concrete layer provides a thin wrapper that casts back to the concrete type: + +```c +// From linear.c +static ag_tensor* nn_linear_module_forward(const nn_module* module, + const ag_tensor* input) +{ + return nn_linear_forward((const nn_linear*)module, input); +} +``` + +This is registered during `_create` via `nn_module_init_base()`. + +### Train/Eval Mode + +`nn_module_set_training` recursively sets the `training` flag on the module and all descendants: + +```c +void nn_module_set_training(nn_module* module, int training) +{ + if (module == NULL) return; + module->training = training != 0; + for (size_t i = 0; i < module->child_count; ++i) { + nn_module_set_training(module->children[i], training); + } +} +``` + +Stochastic modules (like Dropout) check `nn_module_is_training()` during forward. In eval mode, dropout returns the input unchanged. + +### Parameter Counting and Access + +`nn_module_parameter_count` recursively counts all parameters across the entire tree. `nn_module_parameter_at` uses depth-first traversal to access the *i*-th parameter by flat index: + +```c +size_t nn_module_parameter_count(const nn_module* module) +{ + size_t count = module->parameter_count; + for (size_t i = 0; i < module->child_count; ++i) { + count += nn_module_parameter_count(module->children[i]); + } + return count; +} +``` + +### Gradient Zeroing + +`nn_module_zero_grad` recursively zeros gradients on all trainable parameters across the entire tree by calling `ag_zero_grad` on each parameter's `ag_tensor`. + +### Comparison to PyTorch + +| TensorLib | PyTorch | +|---|---| +| `nn_module_register_child()` | `self.child = submodule` (in `__init__`) | +| `nn_module_register_parameter()` | `self.param = nn.Parameter(...)` | +| `nn_module_parameter_count()` | `sum(p.numel() for p in module.parameters())` | +| `nn_module_parameter_at(i)` | `list(module.parameters())[i]` | +| `nn_module_set_training(0)` | `module.eval()` | +| `nn_module_zero_grad()` | `module.zero_grad()` | +| `nn_module_forward()` | `module(input)` | + +--- + +## 3. Parameter System + +### The `nn_parameter` Struct + +```c +struct nn_parameter { + char* name; // unique dotted name, e.g. "proj.weight" + ag_tensor* value; // autograd-wrapped tensor + int trainable; // 1 = accumulate gradients, 0 = frozen +}; +``` + +A parameter wraps an `ag_tensor` (see [autograd_engine.md](./autograd_engine.md)) with a unique name and a trainable flag. The name is used for checkpoint serialization and deduplication. + +```mermaid +flowchart LR + P[nn_parameter] -->|"value"| AT[ag_tensor] + AT -->|"value"| T[tensor] + T -->|"storage"| S[storage] + AT -->|"grad"| G[tensor?] + AT -->|"creator"| C[autograd_node?] + P -->|"name"| N["'proj.weight'"] + P -->|"trainable"| TF["1 / 0"] + style P fill:#e8f4fd,stroke:#2196F3 + style AT fill:#fff3e0,stroke:#FF9800 + style T fill:#e8f5e9,stroke:#4CAF50 + style S fill:#fce4ec,stroke:#E91E63 +``` + +### Creation + +`nn_parameter_create` allocates the tensor, initializes weights according to the chosen policy, and wraps it as an autograd leaf: + +```c +nn_parameter* nn_parameter_create( + const char* name, + int ndim, + const int* dims, + int trainable, + nn_init_kind initializer, + nn_rng* rng +); +``` + +Ownership flows as: caller owns the `nn_parameter*`, the module's `parameters[]` array holds non-owning references. When the module is destroyed via `nn_module_destroy_base`, all registered parameters are destroyed. + +--- + +## 4. Weight Initialization & RNG + +### Splitmix64 PRNG + +TensorLib uses a deterministic [splitmix64](https://xoshiro.di.unimi.it/splitmix64.c) PRNG. The state is a single `uint64_t`: + +```c +struct nn_rng { + uint64_t state; +}; +``` + +The core PRNG produces a 64-bit integer by mixing with the golden ratio constant: + +```c +static uint64_t nn_rng_next_u64(nn_rng* rng) +{ + rng->state += 0x9E3779B97F4A7C15; // golden ratio increment + uint64_t value = rng->state; + value = (value ^ (value >> 30)) * 0xBF58476D1CE4E5B9; + value = (value ^ (value >> 27)) * 0x94D049BB133111EB; + return value ^ (value >> 31); +} +``` + +The top 24 bits are extracted for float conversion, giving ~24-bit precision: + +```c +static float nn_rng_unit(nn_rng* rng) +{ + return (float)(nn_rng_next_u64(rng) >> 40) * (1.0f / 16777216.0f); +} +``` + +### Distributions + +**Uniform:** Simple linear rescaling of the unit random: + +```c +float nn_rng_uniform(nn_rng* rng, float min, float max) +{ + return min + (max - min) * nn_rng_unit(rng); +} +``` + +**Normal:** Box-Muller transform with rejection-sampled `u1` for log safety: + +```c +float nn_rng_normal(nn_rng* rng, float mean, float stddev) +{ + float u1 = ((float)(nn_rng_next_u64(rng) >> 40) + 1.0f) / 16777217.0f; + float u2 = nn_rng_unit(rng); + return mean + stddev * sqrtf(-2.0f * logf(u1)) * cosf(6.28318530717958647692f * u2); +} +``` + +Both return `NAN` on invalid inputs and are fully deterministic given the same seed. + +### Initialization Policies + +The `nn_init_kind` enum controls weight initialization: + +| Policy | Formula | Use Case | +|---|---|---| +| `NN_INIT_ZERO` | `W = 0` | Zero bias, attention QKV init | +| `NN_INIT_ONE` | `W = 1` | LayerNorm weight | +| `NN_INIT_XAVIER_UNIFORM` | `U(-sqrt(6/(fan_in+fan_out)), sqrt(6/(fan_in+fan_out)))` | Linear weights, Embeddings | +| `NN_INIT_XAVIER_NORMAL` | `N(0, sqrt(2/(fan_in+fan_out)))` | Alternative Linear init | +| `NN_INIT_HE_UNIFORM` | `U(-sqrt(6/fan_in), sqrt(6/fan_in))` | ReLU networks | +| `NN_INIT_HE_NORMAL` | `N(0, sqrt(2/fan_in))` | ReLU networks | + +Fan-in and fan-out are computed from the weight tensor shape, treating dimensions beyond the first two as a receptive field multiplier. + +The attention QKV weight uses a custom init of `U(-sqrt(3/C), sqrt(3/C))` applied after zero-init, as a distinct scaling choice. + +--- + +## 5. Individual Layers + +### 5.1 Linear (Fully Connected) + +**Purpose:** Standard affine transformation with optional bias. + +**Parameters:** +- `weight`: shape `[out_features, in_features]` — initialized with `weight_init` +- `bias` (optional): shape `[out_features]` — initialized with `bias_init` + +**Forward pass math:** + +$$y = x W^T + b$$ + +**Implementation:** The weight is transposed via `ag_transpose`, multiplied with the input via `ag_matmul`, and the bias is added via `ag_add`: + +```c +ag_tensor* nn_linear_forward(const nn_linear* layer, const ag_tensor* input) +{ + ag_tensor* transposed = ag_transpose(layer->weight->value, 0, 1); + ag_tensor* product = ag_matmul(input, transposed); + ag_tensor_release(transposed); + if (!layer->use_bias) return product; + ag_tensor* result = ag_add(product, layer->bias->value); + ag_tensor_release(product); + return result; +} +``` + +**Backward gradient:** Computed automatically by the autograd graph — `ag_transpose` and `ag_matmul` each have backward implementations. + +**PyTorch comparison:** Equivalent to `nn.Linear(in_features, out_features, bias=True)`. The weight shape `[out, in]` and the `x @ W^T` convention match PyTorch exactly. + +--- + +### 5.2 Embedding + +**Purpose:** Lookup table that maps integer token IDs to dense vectors. + +**Parameters:** +- `weight`: shape `[vocabulary_size, embedding_width]` — initialized with `weight_init` + +**Forward pass:** Uses `ag_gather_rows` to select rows by token index: + +```c +ag_tensor* nn_embedding_forward(const nn_embedding* layer, const ag_tensor* indices) +{ + return ag_gather_rows(layer->weight->value, indices->value); +} +``` + +Input must be a float tensor containing integer token IDs with `requires_grad == 0`. The output is differentiable w.r.t. the weight matrix. + +**PyTorch comparison:** Equivalent to `nn.Embedding(vocab_size, dim)`. Both store a `[V, D]` weight and look up by row index. Gradients flow into the weight (sparse update). + +--- + +### 5.3 Positional Embedding + +**Purpose:** Adds learned position vectors to token embeddings, giving the model sequence order information. + +**Parameters:** +- Internally owns an `nn_embedding` table of shape `[context_length, embedding_width]` + +**Forward pass:** Generates position IDs `[0, 1, 2, ..., T-1]` from the sequence length, looks them up in the internal embedding table, and adds the result to the token embeddings: + +```c +ag_tensor* nn_positional_embedding_forward( + const nn_positional_embedding* layer, + const ag_tensor* token_embeddings) +{ + // ... generate position_ids [T] ... + positions = nn_embedding_forward(layer->table, position_ids); + result = ag_add(token_embeddings, positions); + return result; +} +``` + +Input shape: `[B, T, C]`. Output shape: `[B, T, C]` (same — addition). + +The internal embedding is registered as a child module, so its parameter appears in the parent's parameter count automatically. + +```mermaid +flowchart TB + TIDS["Token IDs [B, T]"] --> EMB["Token Embedding
Vocab → C"] + EMB --> POS["Positional Embedding
adds learned [T, C] vectors"] + POS --> BLK1["Decoder Block 0"] + BLK1 --> BLK2["Decoder Block 1"] + BLK2 --> BLKN["... Decoder Block N-1"] + BLKN --> NORM["Final LayerNorm"] + NORM --> LM["Language Model Head
Linear [C → V]"] + LM --> LOGITS["Logits [B, T, V]"] + style TIDS fill:#fff3e0 + style LOGITS fill:#e8f5e9 + style EMB fill:#e8f4fd + style LM fill:#fce4ec +``` + +--- + +### 5.4 Layer Normalization + +**Purpose:** Normalizes the last dimension of the input to zero mean and unit variance, with optional learned affine scale and shift. + +**Parameters (when `affine=1`):** +- `weight`: shape `[normalized_width]` — initialized to 1 (`NN_INIT_ONE`) +- `bias`: shape `[normalized_width]` — initialized to 0 (`NN_INIT_ZERO`) + +**Forward pass math:** + +$$\mu = \frac{1}{C} \sum_{i=1}^{C} x_i$$ + +$$\sigma^2 = \frac{1}{C} \sum_{i=1}^{C} (x_i - \mu)^2$$ + +$$\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}$$ + +$$y_i = \gamma_i \hat{x}_i + \beta_i$$ + +**Implementation:** Builds a chain of autograd ops: `ag_mean` → `ag_sub` → `ag_mul` → `ag_mean` → `ag_add` (epsilon) → `ag_sqrt` → `ag_div` → `ag_mul` (gamma) → `ag_add` (beta). + +**PyTorch comparison:** Equivalent to `nn.LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True)`. Uses biased variance (divides by *C*, not *C-1*). + +--- + +### 5.5 Dropout + +**Purpose:** Regularization layer that randomly zeros elements during training. + +**Parameters:** None (stateless — uses the RNG reference). + +**Forward pass math (inverted dropout):** + +$$y_i = \begin{cases} 0 & \text{if } u_i < p \\ \frac{x_i}{1-p} & \text{otherwise} \end{cases}$$ + +where $u_i \sim \text{Uniform}(0, 1)$ and $p$ is the dropout probability. + +**Key behavior:** +- **Training mode:** Samples a binary mask, scales surviving elements by $1/(1-p)$ (inverted dropout). +- **Eval mode:** Returns the input unchanged — no mask sampling, no RNG advancement. +- **Probability = 0:** Returns input directly without allocating a mask. + +```c +if (!nn_module_is_training(&layer->base) || layer->probability == 0.0f) { + ag_tensor_retain((ag_tensor*)input); + return (ag_tensor*)input; +} +``` + +The RNG is stored as a non-owning reference (`nn_dropout` does not own the `nn_rng`). The RNG must outlive the dropout module. + +**PyTorch comparison:** Equivalent to `nn.Dropout(p)`. Both use inverted dropout scaling. The probability must be `[0, 1)`. + +--- + +### 5.6 Multi-Head Attention + +**Purpose:** Decoder-style causal self-attention with fused Q/K/V projections. + +**Parameters:** +- `qkv_weight`: shape `[3, C, C]` — fused Q/K/V projection weights +- `qkv_bias`: shape `[3, C]` — fused Q/K/V biases +- `output`: `nn_linear` child module for the output projection `[C → C]` +- `output_dropout`: `nn_dropout` child module + +**Channels must divide evenly across heads:** `head_width = channels / head_count`. + +```mermaid +flowchart TB + IN["Input [B, T, C]"] --> QKV["Fused QKV Projection
[B, T, 3, C] via [3, C, C]"] + QKV --> Q["Query
[B, T, H, D]"] + QKV --> K["Key
[B, T, H, D]"] + QKV --> V["Value
[B, T, H, D]"] + K --> KT["Key^T
[B, H, D, T]"] + Q --> MATMUL1["Q @ K^T"] + KT --> MATMUL1 + MATMUL1 --> SCALE["÷ sqrt(D)"] + SCALE --> MASK["Causal Mask
+0 or -∞"] + MASK --> SOFTMAX["Softmax"] + SOFTMAX --> MATMUL2["Attn @ V"] + V --> MATMUL2 + MATMUL2 --> TRANSPOSE["Transpose [B, H, T, D] → [B, T, C]"] + TRANSPOSE --> OUT_PROJ["Output Linear [C → C]"] + OUT_PROJ --> DROPOUT["Output Dropout"] + DROPOUT --> OUT["Output [B, T, C]"] + style IN fill:#e8f4fd + style OUT fill:#e8f5e9 + style MASK fill:#fff3e0 + style QKV fill:#fce4ec +``` + +**Forward pass algorithm:** + +1. **Fused QKV projection:** Reshape input to `[1, B, T, C]`, multiply by `[3, 1, C, C]` weight and add `[3, 1, 1, C]` bias. This produces all three projections in one matmul. + +2. **Split and reshape:** Slice along the projection dimension, reshape each to `[B, T, H, D]`, then transpose to `[B, H, T, D]` for batched attention. + +3. **Scaled dot-product attention:** + - Compute scores: $\text{scores} = Q \cdot K^T / \sqrt{D}$ + - Apply causal mask via `nn_apply_causal_mask` (sets positions above the diagonal to $-\infty$) + - Apply softmax: $\text{probs} = \text{softmax}(\text{scores})$ + - Compute context: $\text{context} = \text{probs} \cdot V$ + +4. **Merge heads:** Transpose from `[B, H, T, D]` to `[B, T, C]`. + +5. **Output projection:** Linear layer `[C → C]` followed by dropout. + +**PyTorch comparison:** Similar to `nn.MultiheadAttention(embed_dim, num_heads, dropout)` in self-attention mode, but always causal and with a fused QKV weight matrix. Unlike FlashAttention, TensorLib materializes the full attention matrix — this is a reference implementation, not optimized for memory or compute. + +--- + +### 5.7 MLP (Multi-Layer Perceptron) + +**Purpose:** Configurable feed-forward network with arbitrary hidden sizes and activations. Used as the transformer FFN. + +**Parameters:** Internally creates `hidden_count + 1` `nn_linear` layers registered as children. + +**Configuration:** + +```c +struct nn_mlp_config { + int input_features; + const int* hidden_sizes; // array of hidden layer sizes + size_t hidden_count; // number of hidden layers + int output_features; + const nn_activation* activations; // array of (hidden_count + 1) activations + int use_bias; + nn_init_kind weight_init; + nn_init_kind bias_init; +}; +``` + +**Forward pass:** Chains `Linear → Activation → Linear → Activation → ... → Linear`: + +```c +for (size_t i = 0; i < model->layer_count; ++i) { + current = nn_module_forward(model->base.children[i], current); // Linear + if (model->activations[i].forward != NULL) { + current = model->activations[i].forward(&model->activations[i], current); + } +} +``` + +The final layer uses the last activation from the activations array (typically a no-op or identity for output layers). + +**Built-in activations:** `nn_activation_relu()`, `nn_activation_gelu()`, `nn_activation_sigmoid()`, `nn_activation_tanh()`. Custom activations can be created with `nn_activation_custom()`. + +**PyTorch comparison:** Similar to stacking `nn.Linear` + activations. The decoder block's FFN uses a 2-layer MLP with GELU (hidden size = 4× channels), matching the standard transformer architecture. + +```mermaid +flowchart LR + IN["Input [B, D_in]"] --> L0["Linear [D_in → H_1]"] + L0 --> A0["Activation_0"] + A0 --> L1["Linear [H_1 → H_2]"] + L1 --> A1["Activation_1"] + A1 --> L2["Linear [H_2 → D_out]"] + L2 --> A2["Activation_2"] + A2 --> OUT["Output [B, D_out]"] + style IN fill:#e8f4fd + style OUT fill:#e8f5e9 +``` + +Each `Linear` is registered as a child module; activations are stored in the `activations[]` array and applied between layers. + +--- + +## 6. Loss Functions + +### Softmax and Log-Softmax + +**Numerically stable log-softmax** using the log-sum-exp trick: + +$$\text{log\_softmax}(x_i) = x_i - \max(x) - \log\left(\sum_j e^{x_j - \max(x)}\right)$$ + +**Softmax** is implemented as `exp(log_softmax(x))` for numerical consistency: + +```c +ag_tensor* nn_softmax(const ag_tensor* logits) +{ + ag_tensor* log_probabilities = nn_log_softmax(logits); + ag_tensor* probabilities = ag_exp(log_probabilities); + ag_tensor_release(log_probabilities); + return probabilities; +} +``` + +### Cross-Entropy Loss + +**Implementation:** Computes `NLLLoss(log_softmax(logits), targets)`: + +1. Validates target classes are in `[0, num_classes)` +2. Builds a one-hot selector matrix from the integer targets +3. Computes `log_softmax(logits)` +4. Multiplies with the selector and sums over the class dimension +5. Negates and takes the global mean + +**Formula:** + +$$\mathcal{L} = -\frac{1}{N} \sum_{i=1}^{N} \log p_{y_i}(x_i)$$ + +**PyTorch comparison:** Equivalent to `nn.CrossEntropyLoss()` which combines `nn.LogSoftmax` and `nn.NLLLoss`. The targets are class indices (not one-hot), matching PyTorch's convention. + +### Causal Mask + +`nn_apply_causal_mask` adds a lower-triangular mask to attention scores: + +```c +ag_tensor* nn_apply_causal_mask(const ag_tensor* scores) +{ + // Creates [T, T] mask: 0.0 on and below diagonal, -INFINITY above + for (int row = 0; row < sequence; ++row) + for (int column = 0; column < sequence; ++column) + mask_value->storage->data[row * sequence + column] = + column <= row ? 0.0f : -INFINITY; + return ag_add(scores, mask); +} +``` + +This is applied to 2D slices of the attention score tensor `[B, H, T, T]`. + +--- + +## 7. Optimizers + +### 7.1 SGD + +**Purpose:** Stochastic gradient descent with optional per-parameter learning rates. + +```c +nn_sgd* nn_sgd_create(nn_module* module, float learning_rate); +int nn_sgd_step(nn_sgd* optimizer); +void nn_sgd_zero_grad(nn_sgd* optimizer); +``` + +**Update rule:** + +$$\theta_{t+1} = \theta_t - \eta \cdot g_t$$ + +The SGD optimizer iterates all parameters (recursively via `nn_module_parameter_at`), validates gradient shapes and finiteness, and applies the update element-wise. It skips non-trainable parameters. After updating, it calls `tensor_mark_modified` to invalidate any cached views. + +**Validation:** Before applying updates, `nn_sgd_step` validates: +- Module topology integrity +- Gradient shapes match parameter shapes +- All values and gradients are finite +- The resulting updated values are finite + +### 7.2 AdamW + +**Purpose:** AdamW optimizer with decoupled weight decay, bias correction, and global gradient norm clipping. + +```c +nn_adamw_config nn_adamw_default_config(void); // lr=1e-3, β1=0.9, β2=0.999, ... +nn_adamw* nn_adamw_create(nn_module* module, const nn_adamw_config* config); +int nn_adamw_step(nn_adamw* optimizer); +void nn_adamw_zero_grad(nn_adamw* optimizer); +``` + +**Internal state (per parameter):** +- `first_moments[i]` — exponential moving average of gradients ($m_t$) +- `second_moments[i]` — exponential moving average of squared gradients ($v_t$) +- `steps[i]` — step counter for bias correction + +**Update rules:** + +$$m_t = \beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t$$ + +$$v_t = \beta_2 \cdot v_{t-1} + (1 - \beta_2) \cdot g_t^2$$ + +$$\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}$$ + +$$\theta_{t+1} = \theta_t - \eta \cdot \left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \cdot \theta_t\right)$$ + +**Gradient norm clipping:** If `max_grad_norm > 0`, the global L2 norm of all gradients is computed. If it exceeds `max_grad_norm`, all gradients are scaled down by `max_grad_norm / norm`. + +**PyTorch comparison:** Matches [`torch.optim.AdamW`](https://github.com/pytorch/pytorch/blob/main/torch/optim/adamw.py) with decoupled weight decay (not L2 regularization). Default parameters (`lr=1e-3`, `betas=(0.9, 0.999)`, `eps=1e-8`, `weight_decay=0.01`) match PyTorch's defaults. The gradient clipping is applied inside the step, equivalent to calling `torch.nn.utils.clip_grad_norm_` before the optimizer step. + +### Shared Utilities + +`nn_module_zero_grad` and `nn_clip_grad_norm` are optimizer-independent and available for any module: + +```c +void nn_module_zero_grad(nn_module* module); +int nn_clip_grad_norm(nn_module* module, float max_norm, float* total_norm); +``` + +--- + +## 8. Serialization (Checkpointing) + +TensorLib provides atomic checkpoint save/load for model parameters, AdamW optimizer state, and RNG state. + +### Binary Format + +```mermaid +flowchart LR + subgraph "Checkpoint File" + direction TB + MAGIC["Magic: 'TLCKPT\\0\\0' (8 bytes)"] + VERSION["Version: u32 (1)"] + FLAGS["Flags: u32
bit 0 = AdamW
bit 1 = RNG"] + PCOUNT["Parameter count: u32"] + PARAMS["Parameters (repeated)"] + OPT["AdamW state (if flag set)"] + RNG["RNG state (if flag set)"] + end + subgraph PARAM["Each Parameter"] + direction TB + PNAME["Name: length-prefixed string"] + PNDIM["ndim: u32"] + PDIMS["dims: u32 × ndim"] + PCOUNT2["element count: u64"] + PVALS["values: f32 × count"] + end + subgraph OPTS["AdamW State"] + direction TB + OCONF["Config: 6 × f32 (lr, β1, β2, ε, wd, max_norm)"] + OUCOUNT["parameter count: u32"] + OPER["Per parameter:"] + OPNAME[" name: string"] + OPSTEP[" step: u64"] + OPELcount[" element count: u64"] + OPF1[" first_moments: f32 × count"] + OPF2[" second_moments: f32 × count"] + end + PARAMS --> PARAM + OPT --> OPTS +``` + +### Format Details + +| Field | Type | Description | +|---|---|---| +| Magic | `uint8_t[8]` | `TLCKPT\0\0` | +| Version | `uint32_t` | Currently `1` | +| Flags | `uint32_t` | Bit 0: AdamW state present. Bit 1: RNG state present. | +| Parameter count | `uint32_t` | Number of parameters | +| Parameter name | string | Length-prefixed (u32 length + bytes) | +| Parameter ndim | `uint32_t` | Number of dimensions | +| Parameter dims | `uint32_t[]` | Dimension sizes | +| Element count | `uint64_t` | Total elements (verified against dims) | +| Values | `float[]` | All tensor elements as IEEE 754, little-endian | + +All multi-byte values are serialized in **little-endian** byte order. + +### Atomic Writes + +Save operations are atomic — data is written to a temporary file (`path + ".tmp"`) and then atomically renamed to the final path. On Windows, `MoveFileExA` with `MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH` is used; on POSIX, `rename()`: + +```c +#ifdef _WIN32 + return MoveFileExA(temporary, destination, + MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH) != 0; +#else + return rename(temporary, destination) == 0; +#endif +``` + +If any error occurs during write, the temporary file is deleted and the original file is untouched. + +### Loading Safety + +Loading is **transactional**: the checkpoint is fully parsed and validated against the live module before any state is modified. Validation checks include: +- Magic bytes and version match +- All parameter names are unique and exist in the live module +- All parameter shapes match exactly +- AdamW parameter count and shapes match (if present) +- All values are finite +- No trailing data after the last section + +If any check fails, `nn_checkpoint_load` returns `-1` and leaves all live state intact. + +The loader also supports a legacy format with magic `TLWEIGHT` (no AdamW/RNG sections). + +### PyTorch Comparison + +PyTorch uses Python's `pickle` for `torch.save`/`torch.load`, which can serialize arbitrary Python objects but has security implications. TensorLib's custom binary format is simpler, safer (no arbitrary code execution), and versioned for forward compatibility. However, it only serializes the specific data needed for training (parameters, optimizer state, RNG) — not arbitrary model objects. + +--- + +## 9. API Reference + +### Module System + +| Function | Description | +|---|---| +| `nn_module_forward(module, input)` | Execute forward pass via virtual dispatch | +| `nn_module_set_training(module, training)` | Recursively set train/eval mode | +| `nn_module_is_training(module)` | Query current training flag | +| `nn_module_register_parameter(module, param)` | Register a parameter with the module | +| `nn_module_register_child(module, child)` | Register a sub-module (with cycle detection) | +| `nn_module_parameter_count(module)` | Count all parameters recursively | +| `nn_module_parameter_at(module, index)` | Get parameter by flat index | +| `nn_module_zero_grad(module)` | Recursively zero all parameter gradients | +| `nn_clip_grad_norm(module, max_norm, total_norm)` | Clip gradient L2 norm globally | + +### Activations + +| Function | Description | +|---|---| +| `nn_activation_relu()` | ReLU activation | +| `nn_activation_gelu()` | GELU activation | +| `nn_activation_sigmoid()` | Sigmoid activation | +| `nn_activation_tanh()` | Tanh activation | +| `nn_activation_custom(name, forward, ctx)` | Caller-defined activation | + +### Parameters + +| Function | Description | +|---|---| +| `nn_parameter_create(name, ndim, dims, trainable, init, rng)` | Create a parameter with initialized tensor | +| `nn_parameter_destroy(param)` | Free a parameter and its tensor | + +### RNG + +| Function | Description | +|---|---| +| `nn_rng_seed(rng, seed)` | Set the RNG state | +| `nn_rng_uniform(rng, min, max)` | Sample from Uniform[min, max] | +| `nn_rng_normal(rng, mean, stddev)` | Sample from Normal(mean, stddev) | + +### Layers + +| Function | Description | +|---|---| +| `nn_linear_create(name, in, out, bias, w_init, b_init, rng)` | Create linear layer | +| `nn_linear_forward(layer, input)` | Forward pass: `x @ W^T + b` | +| `nn_linear_destroy(layer)` | Free linear layer | +| `nn_embedding_create(name, vocab, dim, init, rng)` | Create embedding layer | +| `nn_embedding_forward(layer, indices)` | Lookup embeddings by token ID | +| `nn_embedding_destroy(layer)` | Free embedding layer | +| `nn_positional_embedding_create(name, ctx_len, dim, init, rng)` | Create positional embedding | +| `nn_positional_embedding_forward(layer, token_emb)` | Add positional info to token embeddings | +| `nn_positional_embedding_destroy(layer)` | Free positional embedding | +| `nn_layer_norm_create(name, width, eps, affine)` | Create layer normalization | +| `nn_layer_norm_forward(layer, input)` | Normalize last dimension | +| `nn_layer_norm_destroy(layer)` | Free layer normalization | +| `nn_dropout_create(name, probability, rng)` | Create dropout layer | +| `nn_dropout_forward(layer, input)` | Apply dropout (train) or identity (eval) | +| `nn_dropout_destroy(layer)` | Free dropout layer | +| `nn_multihead_attention_create(name, C, heads, drop_p, rng)` | Create multi-head attention | +| `nn_multihead_attention_forward(attn, input)` | Causal self-attention forward | +| `nn_multihead_attention_destroy(attn)` | Free multi-head attention | +| `nn_decoder_block_create(name, C, heads, drop_p, eps, rng)` | Create pre-norm decoder block | +| `nn_decoder_block_forward(block, input)` | Forward through attention + FFN block | +| `nn_decoder_block_destroy(block)` | Free decoder block | +| `nn_decoder_create(name, config, rng)` | Create full decoder (GPT-style) | +| `nn_decoder_forward(decoder, token_ids)` | Forward pass returning logits `[B,T,V]` | +| `nn_decoder_loss(decoder, token_ids, targets)` | Forward + cross-entropy loss (scalar) | +| `nn_decoder_destroy(decoder)` | Free decoder | +| `nn_mlp_create(name, config, rng)` | Create multi-layer perceptron | +| `nn_mlp_forward(model, input)` | Forward through all layers + activations | +| `nn_mlp_destroy(model)` | Free MLP | + +### Losses + +| Function | Description | +|---|---| +| `nn_softmax(logits)` | Numerically stable softmax | +| `nn_log_softmax(logits)` | Numerically stable log-softmax | +| `nn_cross_entropy(logits, targets)` | Cross-entropy loss (mean scalar) | +| `nn_apply_causal_mask(scores)` | Apply lower-triangular causal mask | + +### Optimizers + +| Function | Description | +|---|---| +| `nn_sgd_create(module, lr)` | Create SGD optimizer | +| `nn_sgd_step(optimizer)` | Apply one SGD update step | +| `nn_sgd_zero_grad(optimizer)` | Zero all gradients | +| `nn_sgd_destroy(optimizer)` | Free SGD optimizer | +| `nn_adamw_default_config()` | Get default AdamW config | +| `nn_adamw_create(module, config)` | Create AdamW optimizer | +| `nn_adamw_step(optimizer)` | Apply one AdamW update step | +| `nn_adamw_zero_grad(optimizer)` | Zero all gradients | +| `nn_adamw_destroy(optimizer)` | Free AdamW optimizer | + +### Checkpointing + +| Function | Description | +|---|---| +| `nn_checkpoint_save(path, module, optimizer, rng)` | Save checkpoint (atomic) | +| `nn_checkpoint_load(path, module, optimizer, rng)` | Load checkpoint (transactional) | + +--- + +## 10. Test Coverage + +The test suite consists of **18 test files** covering every component of the NN system: + +| Test File | What It Covers | +|---|---| +| `test_nn_rng.c` | Splitmix64 reproducibility, uniform/normal distribution, boundary checks, determinism across seeds | +| `test_nn_init.c` | All 6 init policies: zero, one, Xavier uniform/normal, He uniform/normal. Verifies reproducibility with same seed, divergence with different seeds, and range/bounds | +| `test_nn_parameter.c` | Parameter creation, name copying (ownership), trainable flag, tensor shape/strides, all init policies, NULL/empty name rejection | +| `test_nn_module.c` | Module creation/destroy, parent/child hierarchy, cycle detection in `register_child`, duplicate rejection, parameter counting, parameter access by index, recursive train/eval propagation, forward dispatch, zero_grad, destroy verification | +| `test_nn_linear.c` | Constructor validation, known-weight forward pass (hand-computed), biasless mode, generic module forward, dimension mismatch rejection | +| `test_nn_embedding.c` | Lookup correctness, gradient rejection (non-differentiable input), shape validation, generic module forward | +| `test_nn_positional_embedding.c` | Exact composition with manual weights, backward pass through positional params, sequence length validation, context length bounds | +| `test_nn_layer_norm.c` | Forward correctness (hand-computed expected values), backward gradient correctness, non-affine mode, epsilon validation, dimension mismatch | +| `test_nn_dropout.c` | Deterministic training (same seed → same mask), inverted dropout scaling, eval mode identity, backward gradient propagation, p=0.0 passthrough, non-training-mode passthrough | +| `test_nn_multihead_attention.c` | Identity QKV projection correctness, causal masking verification, multi-head correctness, dimension mismatch rejection, backward pass | +| `test_nn_decoder_block.c` | Residual connection verification (zero branches → identity), 6-child topology, backward gradient flow, dimension validation | +| `test_nn_decoder.c` | Full topology (all children registered), forward pass shape, loss computation, backward pass, train/eval mode propagation, config validation | +| `test_nn_mlp.c` | Layer count, child registration, forward correctness with ReLU+identity, custom activations, config validation, reproducibility | +| `test_nn_loss.c` | Softmax correctness (with extreme values), log-softmax correctness, cross-entropy forward/backward, target validation | +| `test_nn_causal_mask.c` | Mask correctness (diagonal and above), softmax after masking, backward pass, dimension validation | +| `test_nn_checkpoint.c` | Round-trip save/load, resume training after load, legacy format, flag-only loading (params only, no optimizer/RNG), corruption detection, dimension mismatch, duplicate name rejection, atomic failure cleanup | +| `test_nn_sgd.c` | Recursive update through children, bias/non-trainable skipping, zero_grad, version tracking, invalid module rejection, negative LR rejection | +| `test_nn_adamw.c` | Exact first-step computation (hand-verified), global gradient clipping, zero_grad, config validation, beta bounds, invalid module rejection | + +--- + +## 11. Design Decisions + +### C-Style OOP (Function Pointers) vs C++ Vtables + +TensorLib uses **manual function-pointer dispatch** rather than C++ virtual tables: + +```c +struct nn_module { + nn_module_forward_fn forward; // manual vtable entry + nn_module_destroy_fn destroy; // manual vtable entry +}; +``` + +**Why:** +- C99 is the target language, keeping the project free of C++ complexity and compiler requirements +- Function pointers are explicit and inspectable — no hidden vptr or compiler-injected destructors +- The "first-member embedding" pattern (`nn_linear` has `nn_module base` as its first field) provides safe pointer casting without `container_of` macros, relying on C's guarantee that a pointer to a struct can be cast to a pointer to its first member +- Destroy callbacks are explicit, avoiding C++ destructor ordering surprises + +The trade-off is that adding a new virtual method requires manually wiring it in every layer's `_create` function — there is no compiler enforcement of interface conformance. + +### Manual Memory Management for Modules + +Every `nn_*_create` has a matching `nn_*_destroy`. Modules own their children and parameters: + +- `nn_module_destroy_base` iterates `parameters[]` calling `nn_parameter_destroy`, then iterates `children[]` calling each child's `destroy` callback, and finally frees the arrays +- This means destroying a root module recursively destroys the entire tree + +**Why not reference counting?** +- Simplicity: explicit ownership avoids the complexity of reference cycles and atomic operations +- Performance: no atomic increment/decrement overhead per operation +- Predictability: cleanup is deterministic and happens exactly at `_destroy` calls +- The module tree is acyclic by construction (enforced by `nn_module_register_child`) + +### Pre-Norm vs Post-Norm Transformer + +TensorLib uses **pre-norm** architecture (LayerNorm before attention/FFN, not after): + +```c +// From decoder_block.c +normalized_attention = nn_layer_norm_forward(block->attention_norm, input); +attention_output = nn_multihead_attention_forward(block->attention, normalized_attention); +attention_residual = ag_add(input, attention_output); // residual before FFN +normalized_mlp = nn_layer_norm_forward(block->mlp_norm, attention_residual); +// ... FFN ... +result = ag_add(attention_residual, dropped); +``` + +**Why pre-norm:** +- **Training stability:** Pre-norm transformers are significantly easier to train at depth. The residual stream provides a direct gradient path that bypasses the LayerNorm and attention layers, mitigating vanishing/exploding gradients. +- **No warmup needed:** Pre-norm architectures are less sensitive to learning rate warmup schedules, which is important for a from-scratch implementation without LR scheduling. +- **Industry standard:** GPT-2, GPT-3, LLaMA, and most modern language models use pre-norm. The post-norm variant (original Transformer) requires careful initialization and warmup. + +The trade-off is that pre-norm can be slightly less performant at convergence compared to well-tuned post-norm, but the stability benefits dominate for practical training. diff --git a/docs/tensor_mechanics.md b/docs/tensor_mechanics.md new file mode 100644 index 0000000..41a4287 --- /dev/null +++ b/docs/tensor_mechanics.md @@ -0,0 +1,709 @@ +# Tensor Mechanics + +The tensor component is the foundation of TensorLib. Every higher-level subsystem -- autograd, optimizers, neural-network layers -- operates on `tensor` objects managed by this layer. This document covers the data structures, memory model, view semantics, element-wise and reduction operations, matrix multiplication kernel, and the public API. + +--- + +## Table of Contents + +1. [Overview](#1-overview) +2. [Core Data Structures](#2-core-data-structures) +3. [Memory Management](#3-memory-management) +4. [View Operations](#4-view-operations) +5. [Element-wise Operations](#5-element-wise-operations) +6. [Reductions](#6-reductions) +7. [Matrix Multiplication](#7-matrix-multiplication) +8. [Gather Operations](#8-gather-operations) +9. [API Reference](#9-api-reference) +10. [Test Coverage](#10-test-coverage) +11. [Design Decisions & Tradeoffs](#11-design-decisions--tradeoffs) + +--- + +## 1. Overview + +TensorLib's tensor component provides a fixed-precision (`float32`) n-dimensional array with: + +- **Strided layout** -- any element can be located via `offset + sum(coords[i] * strides[i])`. +- **Zero-copy views** -- reshape, transpose, slice, squeeze, unsqueeze, and expand produce lightweight wrappers that share the underlying storage buffer. +- **Reference counting** -- storage is freed automatically when the last referencing tensor is destroyed. +- **Broadcasting** -- element-wise operations follow right-aligned NumPy-style broadcasting. +- **Hardware-accelerated matmul** -- an AVX2+FMA blocked kernel with packed-RHS optimisation, plus a portable scalar fallback. + +### Comparison with Other Frameworks + +| Aspect | TensorLib | [PyTorch](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/core/TensorBase.h) | [ggml](https://github.com/ggerganov/llama.cpp/blob/master/ggml/include/ggml.h) | NumPy | +|---|---|---|---|---| +| Language | C99 | C++14 (ATen) | C99 | C/Python | +| Storage object | `Storage` with refcount + version | `Storage` with refcount + allocator | Contiguous `float*` in `ggml_tensor` | `ndarray` with base | +| View mechanism | offset + strides on shared `Storage` | `TensorImpl` with `Storage` + `storage_offset` | Stride-based; views rewrite `ne`/`nb` | `ndarray` view with shared buffer | +| Broadcasting | Right-aligned (NumPy) | Right-aligned (NumPy) | Explicit loops | Right-aligned (NumPy) | +| Type system | `float32` only | Multi-dtype (`TensorImpl`) | Per-tensor `ggml_type` | Multi-dtype | + +PyTorch wraps every tensor in a `TensorImpl` that holds a `Storage` pointer, an `IntArrayRef` of sizes/strides, a storage offset, and an `autograd::AutogradMeta` pointer. TensorLib takes a simpler approach: the `tensor` struct directly owns its dimension/stride arrays and offset, and the `Storage` struct carries only the raw data, a reference count, a total element count, and a monotonically increasing version counter. + +ggml's `ggml_tensor` stores dimensions (`ne[4]`) and strides (`nb[4]`) inline in the struct, and the data pointer (`data`) always points to the start of the contiguous allocation. TensorLib uses dynamically allocated `dims`/`strides` arrays of arbitrary rank and an explicit offset field, enabling richer view chains (e.g., slice-then-transpose-then-reshape) without re-materialization. + +--- + +## 2. Core Data Structures + +### `Storage` + +```c +typedef struct { + float* data; // Heap-allocated element buffer + int ref_count; // Number of tensors referencing this storage + int size; // Total number of float elements + uint64_t version; // Monotonic mutation counter (autograd stale-graph detection) +} Storage; +``` + +`Storage` is the shared backing buffer. When a view is created, the new tensor's `storage` pointer is set to the same `Storage` object and `ref_count` is incremented. The `version` field is incremented by `tensor_mark_modified()`; all views share the same version counter, which allows autograd to detect in-place mutations across a view chain. See [autograd_engine.md](./autograd_engine.md) for how this is used in practice. + +### `tensor` + +```c +typedef struct { + Storage* storage; // Shared backing storage (never NULL for a valid tensor) + int ndim; // Number of dimensions (0 = scalar) + int* dims; // Array of dimension sizes (NULL when ndim == 0) + int* strides; // Array of strides in elements (NULL when ndim == 0) + int offset; // Element offset into storage->data +} tensor; +``` + +A tensor is a lightweight descriptor. The `offset` field means a view can start part-way into the storage buffer without any data movement. Strides are measured in *elements* (not bytes), so computing a flat index is: + +```c +int flat = offset; +for (int i = 0; i < ndim; i++) + flat += coords[i] * strides[i]; +``` + +This is implemented in `get_flat_index_nd()` (`tensor_core.c:104`). + +### Structural Diagram + +```mermaid +graph LR + subgraph "tensor A (root)" + TA_STORAGE["storage -->"] + TA_NDIM["ndim = 2"] + TA_DIMS["dims = [2, 3]"] + TA_STRIDES["strides = [3, 1]"] + TA_OFFSET["offset = 0"] + end + + subgraph "tensor B (view: transpose A)" + TB_STORAGE["storage -->"] + TB_NDIM["ndim = 2"] + TB_DIMS["dims = [3, 2]"] + TB_STRIDES["strides = [1, 3]"] + TB_OFFSET["offset = 0"] + end + + subgraph "tensor C (view: slice A)" + TC_STORAGE["storage -->"] + TC_NDIM["ndim = 2"] + TC_DIMS["dims = [1, 3]"] + TC_STRIDES["strides = [3, 1]"] + TC_OFFSET["offset = 3"] + end + + S["Storage\nref_count = 3\nsize = 6\nversion = 0\ndata = [0, 1, 2, 3, 4, 5]"] + + TA_STORAGE --> S + TB_STORAGE --> S + TC_STORAGE --> S +``` + +All three tensors -- the root and its two views -- share a single `Storage`. The transpose view achieves a different logical layout by swapping dims and strides. The slice view adjusts the offset and shrinks one dimension. + +--- + +## 3. Memory Management + +### Allocation + +| Function | Purpose | +|---|---| +| `s_alloc(ndim, dims)` | Allocate a `Storage` with `ref_count = 1`, `version = 0`, and a zero-initialized `float` buffer. | +| `t_alloc(ndim, dims)` | Allocate a `tensor` with its own `Storage`, row-major strides computed via `calc_strides()`, and `offset = 0`. | +| `t_clone(t)` | Deep copy: allocate fresh storage and copy elements in logical order (respects strides/offset). | +| `init_t(c, ref)` | Initialize an already-allocated `tensor` struct to match `ref`'s shape with zero-filled storage. Used internally by autograd node constructors. | + +Zero-sized tensors and negative/zero dimensions are **rejected** by `tensor_checked_numel()`. Scalars are represented as `ndim = 0` with a single-element storage. + +### Reference Counting Lifecycle + +```mermaid +stateDiagram-v2 + [*] --> Allocated : t_alloc() / s_alloc() + Allocated --> Shared : t_transpose() / t_slice() / ...
add_ref_count() bumps ref_count + Shared --> Shared : Another view created
ref_count++ + Shared --> Allocated : View freed
ref_count-- + Allocated --> Freed : Last reference freed
free(data), free(Storage) + Allocated --> Freed : t_free() on sole owner + Freed --> [*] +``` + +When `t_free()` is called on a tensor (`tensor_alloc.c:60`): + +1. If `ref_count > 1`, decrement and free only the tensor metadata (dims, strides). +2. If `ref_count == 1`, free `storage->data`, then the `Storage` struct, then the tensor metadata. + +This means the last tensor to release a shared storage is responsible for freeing the data buffer. + +### Cloning vs. Viewing + +- **`t_clone(t)`** -- allocates a *new* `Storage`, copies every element in logical order (respecting the source's strides/offset), and returns a fully independent tensor. Mutations to the clone do not affect the original, and the clone has its own version counter. +- **View functions** (`t_transpose`, `t_slice`, etc.) -- return a new `tensor` struct that shares the original's `Storage`. The ref_count is incremented. No data is copied. + +```c +tensor* a = t_alloc(2, (int[]){2, 3}); +// a->storage->ref_count == 1 + +tensor* b = t_transpose(a, 0, 1); +// a->storage->ref_count == 2 (a and b share storage) + +t_free(b); +// a->storage->ref_count == 1 (b's release) + +t_free(a); +// storage is freed here +``` + +### Comparison with ggml + +In [ggml](https://github.com/ggerganov/llama.cpp/blob/master/ggml/include/ggml.h), tensors are allocated from a fixed `ggml_context` memory pool. `ggml_new_tensor()` bumps a bump-pointer allocator; there is no reference counting -- the entire context is freed at once. TensorLib's per-tensor reference counting gives more granular lifetime control, which is important for a framework that constructs dynamic computation graphs. + +--- + +## 4. View Operations + +All view operations create a new `tensor` struct that shares the source's `Storage` via `make_view()` (`tensor_view.c:6`). The key parameters manipulated are `dims`, `strides`, and `offset`. + +### `t_transpose(a, dim0, dim1)` + +Swaps the size and stride of two dimensions. No data is moved. + +```c +// Shape [2, 3], strides [3, 1] +tensor* t = t_transpose(a, 0, 1); +// Shape [3, 2], strides [1, 3] -- same storage +``` + +The result is *not* contiguous. See `tensor_view.c:31`. + +### `t_reshape(a, new_ndim, new_dims)` + +Changes the shape while preserving element order. Total element count must match. + +- If `a` is contiguous: returns a zero-copy view with new strides (`tensor_view.c:80`). +- If `a` is not contiguous: clones the data into a fresh contiguous buffer with the new shape (`tensor_view.c:86`). This is the safe default -- PyTorch behaves identically. + +### `t_squeeze(a, dim)` + +Removes a dimension of size 1. The corresponding stride entry is removed from the stride array. + +### `t_unsqueeze(a, dim)` + +Inserts a size-1 dimension at position `dim`. The new stride is computed to maintain contiguous layout: `strides[dim] * dims[dim]` (or `1` at the trailing edge). + +### `t_expand(a, new_ndim, new_dims)` + +Creates a zero-copy view where size-1 dimensions are broadcast by setting their stride to **zero**. This is the same mechanism NumPy and PyTorch use -- a zero stride means every index along that axis reads the same element. + +```c +tensor* a = t_alloc(2, (int[]){2, 1}); // shape [2, 1] +tensor* b = t_expand(a, 2, (int[]){2, 3}); +// b has shape [2, 3], strides [1, 0] +// All three columns of each row read the same value +``` + +Dimensions that already match are kept unchanged. Leading dimensions not present in the source are given stride 0 (implicit size-1 broadcast). See `tensor_view.c:211`. + +### `t_slice(a, dim, start, end)` + +Returns a view into a sub-range along `dim`. The offset is advanced by `start * strides[dim]`, and the size of `dim` becomes `end - start`. All other dimensions and strides are unchanged. + +```c +tensor* a = t_alloc(2, (int[]){3, 4}); +tensor* s = t_slice(a, 0, 1, 3); +// s has shape [2, 4], offset = 1 * 4 = 4 +``` + +### `t_contiguous(t)` + +- If `t` is already contiguous: returns a zero-copy view (shares storage, bumps ref_count). +- If `t` is not contiguous: returns a deep copy in row-major order via `t_clone()`. + +### How Views Share Storage + +```mermaid +graph TD + ROOT["tensor: root\nshape [2, 3, 4]\nstrides [12, 4, 1]\noffset = 0"] + + TR["tensor: transposed\nshape [4, 3, 2]\nstrides [1, 4, 12]\noffset = 0"] + + SL["tensor: sliced\nshape [1, 3, 4]\nstrides [12, 4, 1]\noffset = 12"] + + EX["tensor: expanded\nshape [2, 5, 4]\nstrides [12, 0, 1]\noffset = 0"] + + S["Storage\nref_count = 4\ndata = [...]"] + + ROOT -->|shares| S + TR -->|shares| S + SL -->|shares| S + EX -->|shares| S +``` + +The transpose swaps strides, the slice adjusts offset and shrinks a dimension, and the expand sets a stride to zero -- all without copying data. + +### Comparison with NumPy / PyTorch + +NumPy's `ndarray.view()` creates a view by manipulating strides and offset, exactly like TensorLib. PyTorch's `as_strided()` does the same at the C++ level. TensorLib follows this well-established model, simplified to a single dtype. + +--- + +## 5. Element-wise Operations + +### Binary Operations (Arithmetic) + +| Function | Operator | +|---|---| +| `t_add(a, b)` | `a + b` | +| `t_sub(a, b)` | `a - b` | +| `t_mul(a, b)` | `a * b` | +| `t_div(a, b)` | `a / b` | + +All binary operations: +1. Compute the broadcast output shape via `broadcast_output_shape()` (`tensor_ops.c:42`). +2. Allocate the output tensor. +3. Fast-path: if both inputs and the output have the *same shape and same strides*, use a contiguous SIMD-friendly loop (`add_contiguous`, etc.). +4. Slow-path: iterate over every output element using `advance_coords()` and resolve each input element via `input_index_for_broadcast()` (`tensor_ops.c:76`), which handles size-1 dimension broadcasting by clamping the coordinate to 0. + +**Scalar variants** (`t_add_scalar`, `t_sub_scalar`, `t_mul_scalar`, `t_div_scalar`) avoid allocating a scalar tensor and operate element-wise with a constant float. + +### Broadcasting Rules + +TensorLib uses **right-aligned, NumPy-style broadcasting**. When two operands have different ranks, the shorter tensor is conceptually prepended with dimensions of size 1. Two dimensions are compatible if they are equal or one of them is 1. The output dimension is the maximum of the two. + +```mermaid +graph LR + subgraph "A: shape [2, 3]" + A1["dim 0: 2"] + A2["dim 1: 3"] + end + subgraph "B: shape [1, 3]" + B1["dim 0: 1 (broadcast)"] + B2["dim 1: 3 (match)"] + end + subgraph "Output: shape [2, 3]" + O1["dim 0: 2 (from A)"] + O2["dim 1: 3 (match)"] + end + B1 -->|broadcast| O1 + A1 -->|keep| O1 + B2 -->|match| O2 + A2 -->|match| O2 +``` + +**Example:** `a` has shape `[2, 1]` and `b` has shape `[1, 3]`. The output has shape `[2, 3]`. Each row of `a` is added to each column of `b`. + +This is tested extensively in `test_tensor_ops.c` with cases including rank-0 (scalar) tensors, transposed inputs, and expanded inputs. + +### Unary Operations (Activations) + +| Function | Formula | +|---|---| +| `t_exp(t)` | `exp(x)` | +| `t_log(t)` | `log(x)` | +| `t_relu(t)` | `max(0, x)` | +| `t_tanh(t)` | `tanh(x)` | +| `t_sigmoid(t)` | `1 / (1 + exp(-x))` | +| `t_gelu(t)` | `0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))` | +| `t_pow(t, exp)` | `x^exp` | +| `t_neg(t)` | `-x` | +| `t_sqrt(t)` | `sqrt(x)` | + +All unary operations first materialize the input into a contiguous copy via `t_contiguous()`, apply the function element-wise, and return the output. This avoids branch-heavy strided iteration at the cost of one copy when the input is non-contiguous. IEEE-754 domain results (NaN, -Inf for `log(0)`, NaN for `sqrt(-1)`) are preserved faithfully, as verified by the test suite. + +--- + +## 6. Reductions + +| Function | Behavior | +|---|---| +| `t_sum(a, dim)` | Sum along `dim`, **remove** the reduced axis | +| `t_mean(a, dim)` | Mean along `dim`, **remove** the reduced axis | +| `t_max(a, dim)` | Max along `dim`, **remove** the reduced axis | +| `t_sum_keepdim(a, dim)` | Sum along `dim`, **keep** the reduced axis with size 1 | +| `t_mean_keepdim(a, dim)` | Mean along `dim`, **keep** the reduced axis with size 1 | +| `t_max_keepdim(a, dim)` | Max along `dim`, **keep** the reduced axis with size 1 | + +### Implementation + +Reductions iterate over the output element count. For each output coordinate, the reduced dimension index is set to 0, and a loop accumulates (sum) or compares (max) across all values along that axis (`tensor_reduc.c:33-118` for sum, `tensor_reduc.c:161-258` for max). `t_mean` simply calls `reduce_sum` and divides by the reduction dimension size. + +The `keepdim` variants use `make_reduction_dims()` to compute the output shape, setting the reduced dimension to 1 instead of removing it. + +### Interaction with Autograd + +Reduction gradients broadcast the upstream gradient back to the input shape. For `t_sum`, the gradient is the upstream broadcast to the original shape. For `t_max`, the gradient is scattered only to the positions that held the maximum value. The `keepdim` variants are especially important here -- they produce output shapes that broadcast cleanly against the input shape during the backward pass. See [autograd_engine.md](./autograd_engine.md) for details. + +### IEEE-754 Behavior + +- **`t_sum`**: NaN propagation (any NaN in the reduction window produces NaN in the output). +- **`t_max`**: NaN propagation (if any element is NaN, the result is NaN). `-INFINITY` is handled correctly for all-negative or mixed-infinity inputs. + +--- + +## 7. Matrix Multiplication + +TensorLib provides a multi-level matrix multiplication system, from a portable scalar fallback to a highly optimized AVX2+FMA microkernel. + +### API + +| Function | Description | +|---|---| +| `t_matmul(a, b)` | General matrix multiply with batch broadcasting. Handles vectors, matrices, and batched inputs. | +| `t_pack_matmul_rhs(rhs)` | Pre-pack the RHS into panel format for reuse across multiple matmul calls. | +| `t_matmul_packed_rhs(lhs, rhs)` | Matmul using a pre-packed RHS. | +| `t_free_matmul_packed_rhs(rhs)` | Free the packed RHS buffer. | + +### Kernel Architecture + +The matmul system is implemented in `tensor_matmul.c` and uses a three-level tiling strategy inspired by [OpenBLAS's sgemm microkernel](https://github.com/OpenMathLib/OpenBLAS/blob/master/kernel/x86_64/sgemm_kernel_4x16_haswell.c): + +```mermaid +graph TD + subgraph "Outer Loop: MC x NC blocks" + MC["MC = 64 rows"] + NC["NC = 64 columns"] + end + subgraph "Middle Loop: KC panel" + KC["KC = 128 (inner/accumulation depth)"] + end + subgraph "Micro-kernel: MR x NR tile" + MR["MR = 4 rows"] + NR["NR = 16 columns (2 x 256-bit)"] + end + MC --> KC + NC --> KC + KC --> MR + KC --> NR +``` + +| Parameter | Value | Rationale | +|---|---|---| +| `MR` | 4 | Number of rows processed per micro-kernel invocation. Keeps 4 accumulators (one per row) in YMM registers. | +| `NR` | 16 | Number of columns per micro-kernel. Two 256-bit FMA vectors cover 16 floats. | +| `MC` | 64 | Row block size. Fits L1 cache (64 * 128 * 4B = 32KB for the LHS panel). | +| `NC` | 64 | Column block size. | +| `KC` | 128 | Inner dimension block. Controls packing granularity. | + +### The AVX2+FMA Micro-kernel + +The `matmul_4x16_kernel()` function (`tensor_matmul.c:414`) processes a 4x16 tile of the output: + +```c +// For each k in [0, k_count): +// Load broadcast a[row][k] into __m256 +// Load packed b[k][0:16] into two __m256 +// FMA into 8 accumulators (4 rows x 2 column groups) +for (int k = 0; k < k_count; ++k) { + __m256 b0 = _mm256_loadu_ps(b + 0); + __m256 b1 = _mm256_loadu_ps(b + 8); + __m256 a0_value = _mm256_broadcast_ss(a0 + k); + c00 = _mm256_fmadd_ps(a0_value, b0, c00); + c01 = _mm256_fmadd_ps(a0_value, b1, c01); + // ... rows 1-3 ... +} +``` + +This is an outer-product formulation: each iteration of `k` broadcasts one element of the LHS row and multiplies it against all NR elements of the packed RHS, accumulating into the 4x16 output tile. + +### Packed RHS Optimization + +The `t_pack_matmul_rhs()` function (`tensor_matmul.c:111`) reorganizes the RHS into panels of size `[inner, NR]`, with zero-padding for tail columns. This eliminates stride-guessing in the hot loop and improves cache locality. The packed buffer is 32-byte aligned on Windows (`_aligned_malloc`) for optimal AVX2 loads. + +When `t_matmul()` detects that both operands are contiguous matrices and AVX2 is available, it transparently packs the RHS internally via `try_packed_matrix_matmul()` (`tensor_matmul.c:903`). For non-contiguous operands (e.g., transposed views), it falls back to `matmul_2d_strided()` which handles arbitrary strides. + +The `t_matmul_packed_rhs()` path also packs the LHS block (`pack_lhs_block()`, `tensor_matmul.c:655`) into a temporary buffer of size `MC * KC`, then invokes `matmul_4x16_packed_a_kernel()` -- a variant microkernel that reads both operands from packed buffers. Tail rows/columns (where `rows % MR != 0` or `columns % NR != 0`) are handled by `matmul_2d_packed_rhs_scalar_tails()` (`tensor_matmul.c:671`). + +### Batched and Broadcast Matmul + +`t_matmul()` supports arbitrary batch dimensions. It decomposes each operand into `matmul_operand_info`: + +```c +typedef struct { + int is_vector; // 1D input (promoted to row or column vector) + int batch_rank; // ndim - 2 for matrices, 0 for vectors + int rows; // Last-2 dimension + int inner; // Last-1 dimension (contraction dimension) + int columns; // Last dimension +} matmul_operand_info; +``` + +Batch dimensions are broadcast using the same right-aligned rule as element-wise ops. For each batch coordinate, the appropriate offsets into the LHS, RHS, and output are computed, and a 2D matmul is dispatched. + +**Vector promotion**: 1D inputs are promoted to row vectors (left) or column vectors (right). `vector x vector` produces a scalar (0D output). This matches NumPy/PyTorch semantics. + +### Tiling Strategy Diagram + +```mermaid +graph TD + subgraph "RHS Packing" + RHS_RAW["Raw RHS\nshape [inner, columns]"] + RHS_PACKED["Packed RHS\npanels of [inner, NR]\n32-byte aligned"] + RHS_RAW -->|"pack_rhs_batch()"| RHS_PACKED + end + + subgraph "LHS Packing (packed_rhs path)" + LHS_RAW["Raw LHS (strided)"] + LHS_PACKED["Packed LHS\n[MC, KC] block"] + LHS_RAW -->|"pack_lhs_block()"| LHS_PACKED + end + + subgraph "Micro-kernel" + MK["matmul_4x16_kernel()\n4 rows x 16 cols\n8 x __m256 accumulators\nFMA inner loop"] + end + + LHS_PACKED --> MK + RHS_PACKED --> MK + + OUT["Output Tile\n[MR, NR]"] + MK -->|"accumulate across KC blocks"| OUT +``` + +### Comparison with ggml + +[ggml](https://github.com/ggerganov/llama.cpp/blob/master/ggml/include/ggml.h) dispatches `ggml_mul_mat` through a type-based kernel table (`ggml_compute_forward_mul_mat`), selecting kernels for different quantization types (Q4_0, Q8_0, F16, F32). TensorLib operates exclusively in F32 and selects between AVX2 and scalar at runtime based on CPUID detection (`matmul_avx2_available()`, `tensor_matmul.c:390`). The ggml approach is optimised for quantized inference; TensorLib's approach is optimised for full-precision training. + +### Comparison with OpenBLAS + +[OpenBLAS's Haswell sgemm kernel](https://github.com/OpenMathLib/OpenBLAS/blob/master/kernel/x86_64/sgemm_kernel_4x16_haswell.c) also uses MR=4, NR=16 with FMA inner loops -- the same microkernel dimensions. TensorLib's tiling parameters (MC=64, NC=64, KC=128) are chosen for L1/L2 cache fit on modern x86. OpenBLAS additionally handles prefetch, register tiling across multiple K-blocks, and multi-threading; TensorLib keeps the implementation simpler for now. + +--- + +## 8. Gather Operations + +### `t_gather_rows(table, indices)` + +Selects rows from a rank-2 `table` (shape `[N, W]`) using `indices` of any rank. Indices must be finite, integral float values in `[0, N)`. + +**Output shape**: `indices.shape + [W]` + +```c +tensor* table = t_alloc(2, (int[]){4, 8}); // 4 rows, 8 columns +tensor* idx = t_alloc(2, (int[]){2, 3}); // 2x3 index tensor +tensor* out = t_gather_rows(table, idx); // shape [2, 3, 8] +``` + +This is the building block for embedding lookups, attention mask indexing, and similar operations. The implementation (`tensor_gather.c:20`) iterates over all indices, validates each one (rejecting NaN, Inf, out-of-range, or non-integer values), and copies the corresponding table row into the output. + +--- + +## 9. API Reference + +### Allocation & Lifetime + +| Signature | Description | +|---|---| +| `Storage* s_alloc(int ndim, const int* dims)` | Allocate storage for `prod(dims)` floats, ref_count = 1. | +| `tensor* t_alloc(int ndim, const int* dims)` | Allocate a tensor with its own storage and row-major strides. | +| `void t_free(tensor* t)` | Release a tensor (decrements storage ref_count; frees storage if last reference). NULL-safe. | +| `tensor* t_clone(tensor* t)` | Deep copy; allocates independent storage and copies elements in logical order. | +| `int init_t(tensor* c, tensor* ref)` | Initialize `c` to match `ref`'s shape with zero-filled storage. Returns 0 on success. | +| `void add_ref_count(Storage* a, tensor* b)` | Link `b` to storage `a` and increment ref_count. | +| `void tensor_mark_modified(tensor* value)` | Increment storage version counter (autograd mutation detection). | + +### Core Utilities + +| Signature | Description | +|---|---| +| `int tensor_numel(tensor* t)` | Return total element count, or 0 for invalid tensors. | +| `int is_contiguous(tensor* t)` | Returns 1 if strides match row-major layout. | +| `void calc_strides(int ndim, const int* dims, int* strides)` | Compute row-major strides. | +| `int get_flat_index_nd(tensor* t, int* coords)` | Compute flat storage index from n-dimensional coordinates. | +| `int same_shape(tensor* a, tensor* b)` | 1 if both have identical ndim and dims. | +| `int same_stride(tensor* a, tensor* b)` | 1 if both have identical ndim and strides. | +| `void advance_coords(int* coords, const int* dims, int ndim)` | Increment n-dimensional coordinate by one position (row-major order). | +| `int tensor_has_valid_shape(const tensor* t)` | Validate dimensions are positive and non-overflowing. | +| `int tensor_has_valid_layout(const tensor* t)` | Validate shape + non-negative strides/offset. | +| `int tensor_has_valid_metadata(const tensor* t)` | Full validation: layout + storage bounds check. | +| `int tensor_checked_numel(int ndim, const int* dims, size_t* result)` | Compute numel with overflow and positivity checks. | +| `int tensor_copy_metadata(int ndim, const int* dims, const int* strides, int** out_dims, int** out_strides)` | Deep-copy dims and strides arrays. Returns 0 on success. | + +### View Operations + +| Signature | Description | +|---|---| +| `tensor* t_transpose(tensor* a, int dim0, int dim1)` | Swap two dimensions (zero-copy view). | +| `tensor* t_reshape(tensor* a, int new_ndim, int* new_dims)` | Change shape; view if contiguous, clone otherwise. | +| `tensor* t_squeeze(tensor* a, int dim)` | Remove a size-1 dimension. | +| `tensor* t_unsqueeze(tensor* a, int dim)` | Insert a size-1 dimension at `dim`. | +| `tensor* t_expand(tensor* a, int new_ndim, const int* new_dims)` | Broadcast via zero strides. | +| `tensor* t_slice(tensor* a, int dim, int start, int end)` | Sub-range view along `dim`. | +| `tensor* t_contiguous(tensor* t)` | View if already contiguous, clone otherwise. | + +### Element-wise Operations + +| Signature | Description | +|---|---| +| `tensor* t_add(tensor* a, tensor* b)` | Element-wise addition with broadcasting. | +| `tensor* t_sub(tensor* a, tensor* b)` | Element-wise subtraction with broadcasting. | +| `tensor* t_mul(tensor* a, tensor* b)` | Element-wise multiplication with broadcasting. | +| `tensor* t_div(tensor* a, tensor* b)` | Element-wise division with broadcasting. | +| `tensor* t_add_scalar(tensor* a, float scalar)` | Add scalar to every element. | +| `tensor* t_sub_scalar(tensor* a, float scalar)` | Subtract scalar from every element. | +| `tensor* t_mul_scalar(tensor* a, float scalar)` | Multiply every element by scalar. | +| `tensor* t_div_scalar(tensor* a, float scalar)` | Divide every element by scalar. | +| `tensor* t_exp(tensor* t)` | Element-wise exp. | +| `tensor* t_log(tensor* t)` | Element-wise log. | +| `tensor* t_relu(tensor* t)` | Element-wise ReLU. | +| `tensor* t_tanh(tensor* t)` | Element-wise tanh. | +| `tensor* t_sigmoid(tensor* t)` | Element-wise sigmoid. | +| `tensor* t_gelu(tensor* t)` | Element-wise GELU (tanh approximation). | +| `tensor* t_pow(tensor* t, float exponent)` | Element-wise power. | +| `tensor* t_neg(tensor* t)` | Element-wise negation. | +| `tensor* t_sqrt(tensor* t)` | Element-wise square root. | + +### Reductions + +| Signature | Description | +|---|---| +| `tensor* t_sum(tensor* a, int dim)` | Sum along axis, remove it. | +| `tensor* t_mean(tensor* a, int dim)` | Mean along axis, remove it. | +| `tensor* t_max(tensor* a, int dim)` | Max along axis, remove it. | +| `tensor* t_sum_keepdim(tensor* a, int dim)` | Sum along axis, keep size-1. | +| `tensor* t_mean_keepdim(tensor* a, int dim)` | Mean along axis, keep size-1. | +| `tensor* t_max_keepdim(tensor* a, int dim)` | Max along axis, keep size-1. | + +### Matrix Multiplication + +| Signature | Description | +|---|---| +| `tensor* t_matmul(tensor* a, tensor* b)` | General matmul with batch broadcasting and vector promotion. | +| `tensor_matmul_packed_rhs* t_pack_matmul_rhs(const tensor* rhs)` | Pre-pack RHS into panel format. | +| `tensor* t_matmul_packed_rhs(const tensor* lhs, const tensor_matmul_packed_rhs* rhs)` | Matmul with pre-packed RHS. | +| `void t_free_matmul_packed_rhs(tensor_matmul_packed_rhs* rhs)` | Free packed RHS buffer. | + +### Gather + +| Signature | Description | +|---|---| +| `tensor* t_gather_rows(tensor* table, tensor* indices)` | Row-lookup from a rank-2 table. Output shape is `indices.shape + [table_width]`. | + +--- + +## 10. Test Coverage + +Each source file in `src/tensor/` has a corresponding test file in `tests/unit/tensor/`. Tests use a lightweight `TEST()`/`RUN_TEST()`/`ASSERT_*` macro framework (defined in `tests/fixtures/test_common.h`). + +### `test_tensor_alloc.c` (17 tests) + +Covers: +- `s_alloc` basic allocation, ndim-0 scalar, null-dims rejection, non-positive and overflowing dimensions +- `t_alloc` shape/strides correctness, scalar allocation, negative ndim rejection +- `t_free` null safety, shared-storage ref_count decrement +- `init_t` shape copy, zero-fill, null-argument safety +- `t_clone` deep copy independence, strided-view materialization, version counter independence +- `add_ref_count` linking, null safety + +### `test_tensor_core.c` (18 tests) + +Covers: +- `calc_strides` for 1D, 3D, and ndim-0 (no-op) +- `advance_coords` normal increment, carry, and wrap-around +- `get_flat_index_nd` contiguous and strided lookup, null safety +- `same_shape` and `same_stride` true/false cases, null args +- `is_contiguous` row-major true, transposed false, null +- `tensor_numel` rejection of negative, zero, and overflowing shapes + +### `test_tensor_view.c` (18 tests) + +Covers: +- `t_transpose` dims/strides swap, storage sharing, ref_count, contiguity loss, bounds checking, null input +- `t_contiguous` on contiguous (view returned) and transposed (clone), null +- `t_reshape` contiguous input (view), strided input (clone), element count mismatch, invalid/overflowing dims +- `t_slice` offset/strides correctness, invalid arguments +- `t_unsqueeze` dimension insertion, stride computation, axis bounds +- `t_expand` size-1 broadcasting, stride=0, leading dimensions, incompatible shapes + +### `test_tensor_ops.c` (15 tests) + +Covers: +- `t_add` contiguous, independent storage, strided (transposed) input, shape mismatch +- Broadcasting: singleton dimensions, strided + expanded inputs, rank-0 scalar tensors +- Scalar helpers: `t_add_scalar`, `t_sub_scalar`, `t_mul_scalar`, `t_div_scalar` +- `t_sub`, `t_mul`, `t_div` with IEEE edge cases (NaN, +Inf, 0/0) +- All binary ops with transposed inputs (logical order verification) +- Unary ops (`t_neg`, `t_sqrt`, `t_exp`, `t_log`, `t_relu`, `t_gelu`, `t_sigmoid`, `t_tanh`, `t_pow`): contiguous values, transposed inputs, IEEE-754 domain results +- Null-input rejection for all operations + +### `test_tensor_reduc.c` (21 tests) + +Covers: +- `t_sum` every axis of a 3D tensor, 1D to scalar, unit reduction dimension +- `t_sum` on transposed, sliced, reshaped, unsqueezed, squeezed, and expanded views +- `t_sum` IEEE-754 NaN propagation +- `t_mean` floating-point division, strided views, null/invalid args +- `t_max` every axis, 1D with negative values, unit dimension, transposed/offset views, expanded broadcast +- `t_max` NaN propagation, infinity handling (-Inf, +Inf, all-negative-Inf) +- `keepdim` variants for sum, mean, max: output shape preservation, views, rank-1 inputs + +### `test_tensor_matmul.c` (15 tests) + +Covers: +- 2D matmul correctness +- 3D batched matmul +- Batch dimension broadcasting (4D with singleton dims) +- Transposed views as inputs +- Vector cases: dot product (1D x 1D -> 0D), matrix-vector, vector-matrix +- Packed RHS: snapshot of transposed view, transposed slices with kernel tails, non-multiple kernel dimensions (5x129 * 129x17), reshape/squeeze/unsqueeze chains, batch broadcasting, materialized contiguous views +- Packed RHS rejection of zero-stride (expanded) and non-matrix RHS +- Invalid shape rejection (mismatched inner dims, incompatible batches, scalars) +- Vector inner dimension validation +- Aliased input (self-matmul) with independent output storage + +### Key Test Patterns + +1. **Shape and value verification**: Allocate, fill, operate, check ndim, dims, and every element. +2. **Storage independence**: Assert `c->storage != a->storage` after operations. +3. **View chain testing**: Transpose-then-slice-then-reshape chains to verify offset/stride correctness. +4. **IEEE-754 domain checks**: NaN propagation, infinity handling, `0.0/0.0 = NaN`, `x/0.0 = +Inf`. +5. **Null and boundary rejection**: Pass NULL, negative dims, zero dims, overflow dims, out-of-range axes. + +--- + +## 11. Design Decisions & Tradeoffs + +### Why C99? + +C99 provides the minimal runtime and maximum portability needed for a tensor library that may be embedded in larger C/C++ projects. Features like `restrict`, `stdint.h`, `stdbool.h`, compound literals, and variable-length arrays (used sparingly) are sufficient. The choice avoids C++ ABI complications and makes the library trivially linkable from any language with a C FFI. + +### Why Not Use BLAS? + +A built-in matmul implementation gives full control over: + +1. **View-aware dispatch** -- the kernel can operate directly on strided inputs without requiring the caller to materialize contiguous copies. BLAS routines like `sgemm` require leading-dimension parameters that don't map cleanly to arbitrary stride patterns. +2. **Batch integration** -- the batch loop is tightly integrated with the kernel dispatch, avoiding per-batch allocations. +3. **Packed RHS lifecycle** -- the packed buffer format is owned by TensorLib and managed through its own API, enabling the `t_pack_matmul_rhs` / `t_matmul_packed_rhs` workflow. +4. **Deployment simplicity** -- no dependency on external BLAS libraries, which simplifies linking on Windows, embedded systems, and CI environments. + +The AVX2+FMA kernel achieves competitive performance for the matrix sizes typical in deep learning workloads. For very large matrices, users can swap in a BLAS-backed implementation without changing the public API. + +### Storage Version Counter for Stale-Graph Detection + +The `version` field in `Storage` is incremented by `tensor_mark_modified()`. Since all views share the same storage, modifying one view bumps the version visible to all others. This enables the autograd engine to detect in-place mutations that would invalidate cached intermediate values in the computation graph. See [autograd_engine.md](./autograd_engine.md). + +### Why Right-Aligned Broadcasting? + +Right-aligned (trailing-dimension) broadcasting is the convention established by NumPy and adopted by PyTorch, TensorFlow, and JAX. It aligns the *innermost* dimensions, which correspond to the most frequently computed axes (spatial, channel, feature). Right-aligned broadcasting means the most common patterns -- adding a bias vector to a batch of activations, scaling a matrix by a scalar -- require no transposition or shape manipulation by the caller. + +### Element Count Validation + +TensorLib intentionally rejects tensors with zero-sized dimensions or total element count exceeding `INT_MAX`. This avoids subtle bugs where zero-sized tensors create empty storage objects (which can cause division-by-zero in reductions or degenerate pointer arithmetic in kernels), and ensures all dimension and offset arithmetic stays within `int` range for predictable behavior.