topos 0.13.1

//! GPT-2 (124M) as a module tree: the released checkpoint's layout,
//! expressed in topos's composition tier.
//!
//! Every struct here is an ordinary [`Module`] implementation — the
//! blocks are structs of [`Linear`]s and [`LayerNorm`]s stacked in a
//! plain `Vec`, and attention is a custom module built from two
//! `Linear`s and the public op surface. The tree's `visit` paths
//! mirror the checkpoint's own tensor names (`h.{i}.attn.c_attn`,
//! `ln_f`, ...), so loading the pretrained weights is one
//! [`named_restore`] over the paths the model announces itself; the
//! adapter shrinks to the leaf spellings the module tier and the
//! checkpoint disagree on.
//!
//! Beside the full-context `express`, the tree records a one-token
//! decode step (`notes/carry.md`): each layer's keys and values live
//! in caller-carried capacity-shaped caches, the step appends the new
//! token's rows by a `scatter` over the position one-hot, and only
//! one stream row flows through the stack. The two expressions share
//! the same parameters on the same tape, so one `Parameters` value
//! serves both plans.
//!
//! The tree is generic over the element type: the same structs record
//! the f32 model and the `Bf16` one, which is the genericity the
//! module design promises.

use topos::checkpoint::named_restore;
use topos::{
    Element, LayerNorm, Linear, Module, Parameters, Path, Segment, Symbol, Tape, Tensor, Value,
    Visitor, concat, named_parameters, scaled_dot_product,
};

use crate::weights::Weights;

/// How many tokens of context the recorded graph attends over.
pub const CONTEXT_LEN: usize = 256;

/// How many dimensions the residual stream has.
pub const EMBED_DIM: usize = 768;

/// How many attention heads split the stream.
const HEAD_COUNT: usize = 12;

/// How many dimensions each head reads and writes.
const HEAD_DIM: usize = EMBED_DIM / HEAD_COUNT;

/// How many transformer blocks the model stacks.
const LAYER_COUNT: usize = 12;

/// How many tokens the vocabulary holds.
pub const VOCABULARY_LEN: usize = 50257;

/// How many positions the released position table covers.
const POSITION_COUNT: usize = 1024;

/// The GELU tanh approximation the checkpoint was trained with, its
/// constants held as scalar leaves shared by every block — float
/// constants are caller territory, so the facade tier cannot supply
/// this activation.
#[derive(Clone)]
struct Gelu {
    half: Symbol,
    one: Symbol,
    root: Symbol,
    coefficient: Symbol,
}

impl Gelu {
    fn new<E: Element + From<f32>>(tape: &Tape<E>) -> Self {
        let scalar = |value: f32| tape.leaf(Tensor::filled([], E::from(value))).symbol();
        Self {
            half: scalar(0.5),
            one: scalar(1.0),
            // The square root of 2 over pi, as the checkpoint's
            // training defined it.
            root: scalar(0.797_884_6),
            coefficient: scalar(0.044_715),
        }
    }
}

impl<E: Element> Module<E> for Gelu {
    /// Records `0.5 x (1 + tanh(sqrt(2/pi) (x + 0.044715 x^3)))`.
    ///
    /// The constants are leaves, not parameters, so the default no-op
    /// `visit` is right: the checkpoint has nothing to restore here.
    fn express<'tape>(&self, input: Value<'tape, E>) -> Value<'tape, E> {
        let tape = input.tape();
        let half = tape.resolve(self.half);
        let one = tape.resolve(self.one);
        let root = tape.resolve(self.root);
        let coefficient = tape.resolve(self.coefficient);
        let cubic = input * input * input * coefficient.broadcast_like(input);
        let inner = ((input + cubic) * root.broadcast_like(input)).tanh();
        input * (inner + one.broadcast_like(inner)) * half.broadcast_like(input)
    }
}

/// One decode step's results: the output row and the caches with the
/// step's key and value rows appended.
struct Decoded<'tape, E> {
    output: Value<'tape, E>,
    keys: Value<'tape, E>,
    values: Value<'tape, E>,
}

/// Multi-head causal self-attention over the fused query-key-value
/// projection the checkpoint ships: every head is a rank-2 slice of
/// one `c_attn` output, and the heads concatenate back through
/// `c_proj`.
struct Attention<E> {
    fused: Linear<E>,
    projection: Linear<E>,
    mask: Symbol,
    scale: Symbol,
}

impl<E: Element + From<f32>> Attention<E> {
    /// Allocates the projections with placeholder payloads; `mask` and
    /// `scale` are leaves shared by every block.
    fn new(tape: &Tape<E>, mask: Symbol, scale: Symbol) -> Self {
        let zeros = |shape: [usize; 2]| Tensor::filled(shape, E::from(0.0));
        let bias = |extent: usize| Tensor::filled([extent], E::from(0.0));
        Self {
            fused: Linear::new(tape, zeros([EMBED_DIM, 3 * EMBED_DIM]), bias(3 * EMBED_DIM)),
            projection: Linear::new(tape, zeros([EMBED_DIM, EMBED_DIM]), bias(EMBED_DIM)),
            mask,
            scale,
        }
    }
}

impl<E: Element> Attention<E> {
    /// Records the one-token decode step: the row's fused projection
    /// appends its key and value rows into the caches (a `scatter`
    /// by the position one-hot over a still-zero row, so the append
    /// is a pure add), and the row's query attends over the whole
    /// capacity under the fed mask row. The full-context `mask` leaf
    /// plays no part here.
    fn express_decode<'tape>(
        &self,
        tape: &'tape Tape<E>,
        input: Value<'tape, E>,
        keys: Value<'tape, E>,
        values: Value<'tape, E>,
        position: Value<'tape, E>,
        mask: Value<'tape, E>,
    ) -> Decoded<'tape, E> {
        let scale = tape.resolve(self.scale);
        let fused = self.fused.express(input);
        let keys = keys + fused.narrow(1, EMBED_DIM, EMBED_DIM).scatter(position);
        let values = values + fused.narrow(1, 2 * EMBED_DIM, EMBED_DIM).scatter(position);
        let heads: Vec<Value<'tape, E>> = (0..HEAD_COUNT)
            .map(|head| {
                let query = fused.narrow(1, head * HEAD_DIM, HEAD_DIM);
                let key = keys.narrow(1, head * HEAD_DIM, HEAD_DIM);
                let value = values.narrow(1, head * HEAD_DIM, HEAD_DIM);
                scaled_dot_product(query, key, value, mask, scale)
            })
            .collect();
        Decoded {
            output: self.projection.express(concat(&heads, 1)),
            keys,
            values,
        }
    }
}

impl<E: Element> Module<E> for Attention<E> {
    fn express<'tape>(&self, input: Value<'tape, E>) -> Value<'tape, E> {
        let tape = input.tape();
        let mask = tape.resolve(self.mask);
        let scale = tape.resolve(self.scale);
        let fused = self.fused.express(input);
        let heads: Vec<Value<'tape, E>> = (0..HEAD_COUNT)
            .map(|head| {
                let query = fused.narrow(1, head * HEAD_DIM, HEAD_DIM);
                let key = fused.narrow(1, EMBED_DIM + head * HEAD_DIM, HEAD_DIM);
                let value = fused.narrow(1, 2 * EMBED_DIM + head * HEAD_DIM, HEAD_DIM);
                scaled_dot_product(query, key, value, mask, scale)
            })
            .collect();
        self.projection.express(concat(&heads, 1))
    }

    fn visit(&self, visitor: &mut dyn Visitor) {
        visitor.enter(Segment::Name("c_attn"));
        self.fused.visit(visitor);
        visitor.leave();
        visitor.enter(Segment::Name("c_proj"));
        self.projection.visit(visitor);
        visitor.leave();
    }
}

/// The block's GELU MLP: up projection, activation, down projection.
struct FeedForward<E> {
    up: Linear<E>,
    activation: Gelu,
    down: Linear<E>,
}

impl<E: Element + From<f32>> FeedForward<E> {
    fn new(tape: &Tape<E>, activation: Gelu) -> Self {
        let zeros = |shape: [usize; 2]| Tensor::filled(shape, E::from(0.0));
        let bias = |extent: usize| Tensor::filled([extent], E::from(0.0));
        Self {
            up: Linear::new(tape, zeros([EMBED_DIM, 4 * EMBED_DIM]), bias(4 * EMBED_DIM)),
            activation,
            down: Linear::new(tape, zeros([4 * EMBED_DIM, EMBED_DIM]), bias(EMBED_DIM)),
        }
    }
}

impl<E: Element> Module<E> for FeedForward<E> {
    fn express<'tape>(&self, input: Value<'tape, E>) -> Value<'tape, E> {
        let lifted = self.up.express(input);
        let hidden = self.activation.express(lifted);
        self.down.express(hidden)
    }

    fn visit(&self, visitor: &mut dyn Visitor) {
        visitor.enter(Segment::Name("c_fc"));
        self.up.visit(visitor);
        visitor.leave();
        visitor.enter(Segment::Name("c_proj"));
        self.down.visit(visitor);
        visitor.leave();
    }
}

/// One pre-norm transformer block: attention and the MLP each read
/// their own normalization of the stream and add back into it.
struct Block<E> {
    attention_norm: LayerNorm<E>,
    attention: Attention<E>,
    hidden_norm: LayerNorm<E>,
    feed_forward: FeedForward<E>,
}

impl<E: Element + From<f32>> Block<E> {
    fn new(tape: &Tape<E>, mask: Symbol, scale: Symbol, activation: Gelu) -> Self {
        Self {
            attention_norm: layer_norm(tape),
            attention: Attention::new(tape, mask, scale),
            hidden_norm: layer_norm(tape),
            feed_forward: FeedForward::new(tape, activation),
        }
    }
}

impl<E: Element> Block<E> {
    /// Records the block's one-token decode step over the stream row:
    /// the same pre-norm wiring as `express`, with attention reading
    /// and updating the layer's caches.
    fn express_decode<'tape>(
        &self,
        tape: &'tape Tape<E>,
        input: Value<'tape, E>,
        keys: Value<'tape, E>,
        values: Value<'tape, E>,
        position: Value<'tape, E>,
        mask: Value<'tape, E>,
    ) -> Decoded<'tape, E> {
        let attended = self.attention.express_decode(
            tape,
            self.attention_norm.express(input),
            keys,
            values,
            position,
            mask,
        );
        let stream = input + attended.output;
        let lifted = self.feed_forward.express(self.hidden_norm.express(stream));
        Decoded {
            output: stream + lifted,
            keys: attended.keys,
            values: attended.values,
        }
    }
}

impl<E: Element> Module<E> for Block<E> {
    fn express<'tape>(&self, input: Value<'tape, E>) -> Value<'tape, E> {
        let attended = self.attention.express(self.attention_norm.express(input));
        let stream = input + attended;
        let lifted = self.feed_forward.express(self.hidden_norm.express(stream));
        stream + lifted
    }

    fn visit(&self, visitor: &mut dyn Visitor) {
        visitor.enter(Segment::Name("ln_1"));
        self.attention_norm.visit(visitor);
        visitor.leave();
        visitor.enter(Segment::Name("attn"));
        self.attention.visit(visitor);
        visitor.leave();
        visitor.enter(Segment::Name("ln_2"));
        self.hidden_norm.visit(visitor);
        visitor.leave();
        visitor.enter(Segment::Name("mlp"));
        self.feed_forward.visit(visitor);
        visitor.leave();
    }
}

/// Builds a layer norm with the conventional placeholder payloads and
/// the epsilon the checkpoint was trained with.
fn layer_norm<E: Element + From<f32>>(tape: &Tape<E>) -> LayerNorm<E> {
    LayerNorm::new(
        tape,
        Tensor::filled([EMBED_DIM], E::from(1.0)),
        Tensor::filled([EMBED_DIM], E::from(0.0)),
        Tensor::filled([], E::from(1e-5)),
    )
}

/// The whole model: token and position tables, the block stack, and
/// the final norm. The language-model head is not a child — it is the
/// embedding table tied, recorded transposed by the caller through
/// [`Gpt2::embeddings`].
pub struct Gpt2<E> {
    embeddings: Symbol,
    positions: Symbol,
    blocks: Vec<Block<E>>,
    final_norm: LayerNorm<E>,
}

impl<E: Element + From<f32> + 'static> Gpt2<E> {
    /// Allocates the model's parameters with placeholder payloads, in
    /// visit order.
    ///
    /// Construction order is a contract: the emitted plan's leading
    /// arguments are the parameters in recording order, so recording
    /// them in visit order makes the positional snapshot exactly the
    /// emitted argument list.
    pub fn new(tape: &Tape<E>) -> Self {
        let embeddings = tape
            .parameter(Tensor::filled([VOCABULARY_LEN, EMBED_DIM], E::from(0.0)))
            .symbol();
        let positions = tape
            .parameter(Tensor::filled([POSITION_COUNT, EMBED_DIM], E::from(0.0)))
            .symbol();

        // The causal mask and the head scale are leaves shared by all
        // twelve blocks, like the GELU constants; leaves embed in the
        // plan as constants, so none of them join the argument list.
        let mask_elements: Vec<E> = (0..CONTEXT_LEN * CONTEXT_LEN)
            .map(|at| {
                if at % CONTEXT_LEN <= at / CONTEXT_LEN {
                    E::from(0.0)
                } else {
                    E::from(f32::NEG_INFINITY)
                }
            })
            .collect();
        let mask = tape
            .leaf(Tensor::new([CONTEXT_LEN, CONTEXT_LEN], mask_elements))
            .symbol();
        let scale = tape
            .leaf(Tensor::filled([], E::from(1.0 / (HEAD_DIM as f32).sqrt())))
            .symbol();
        let activation = Gelu::new(tape);

        let blocks = (0..LAYER_COUNT)
            .map(|_| Block::new(tape, mask, scale, activation.clone()))
            .collect();
        Self {
            embeddings,
            positions,
            blocks,
            final_norm: layer_norm(tape),
        }
    }

    /// Returns the symbol of the `[vocabulary, embed]` token table:
    /// the typed accessor the tied language-model head reads.
    pub fn embeddings(&self) -> Symbol {
        self.embeddings
    }
}

impl<E: Element> Gpt2<E> {
    /// Returns how many blocks the model stacks: one cache pair per
    /// block in the decode step.
    pub fn layers(&self) -> usize {
        self.blocks.len()
    }

    /// Records the one-token decode step over the embedded `[1, embed]`
    /// row: the position row (gathered by the same one-hot that places
    /// the cache appends), the block stack over the caches, the final
    /// norm.
    ///
    /// `caches` pairs each layer's key and value cache inputs in block
    /// order; the returned pairs are the updated caches in the same
    /// order.
    pub fn express_decode<'tape>(
        &self,
        tape: &'tape Tape<E>,
        embedded: Value<'tape, E>,
        caches: &[(Value<'tape, E>, Value<'tape, E>)],
        position: Value<'tape, E>,
        mask: Value<'tape, E>,
    ) -> (Value<'tape, E>, Vec<(Value<'tape, E>, Value<'tape, E>)>) {
        assert_eq!(caches.len(), self.blocks.len(), "one cache pair per block");
        let positions = tape.resolve(self.positions);
        let row = positions.narrow(0, 0, CONTEXT_LEN).gather(position);
        let mut stream = embedded + row;
        let mut updated = Vec::with_capacity(caches.len());
        for (block, &(keys, values)) in self.blocks.iter().zip(caches) {
            let decoded = block.express_decode(tape, stream, keys, values, position, mask);
            stream = decoded.output;
            updated.push((decoded.keys, decoded.values));
        }
        (self.final_norm.express(stream), updated)
    }
}

impl<E: Element> Module<E> for Gpt2<E> {
    /// Records the model over the embedded `[context, embed]` window:
    /// positions in, the block stack, the final norm.
    fn express<'tape>(&self, input: Value<'tape, E>) -> Value<'tape, E> {
        let tape = input.tape();
        let positions = tape.resolve(self.positions);
        let context = input.shape().axes()[0];
        let stream = input + positions.narrow(0, 0, context);
        let stream = self
            .blocks
            .iter()
            .fold(stream, |value, block| block.express(value));
        self.final_norm.express(stream)
    }

    fn visit(&self, visitor: &mut dyn Visitor) {
        visitor.enter(Segment::Name("wte"));
        visitor.parameter("weights", self.embeddings);
        visitor.leave();
        visitor.enter(Segment::Name("wpe"));
        visitor.parameter("weights", self.positions);
        visitor.leave();
        visitor.enter(Segment::Name("h"));
        for (index, block) in self.blocks.iter().enumerate() {
            visitor.enter(Segment::Index(index));
            block.visit(visitor);
            visitor.leave();
        }
        visitor.leave();
        visitor.enter(Segment::Name("ln_f"));
        self.final_norm.visit(visitor);
        visitor.leave();
    }
}

/// Renders `path` as the checkpoint's tensor name. The tree mirrors
/// the released layout, so only the leaf spellings differ: the module
/// tier's `weights` and `scale` are the checkpoint's `weight`, and
/// its `shift` is the checkpoint's `bias`.
fn foreign_name(path: &Path) -> String {
    let segments = path.segments();
    let mut name = String::new();
    for (position, segment) in segments.iter().enumerate() {
        if position > 0 {
            name.push('.');
        }
        if position + 1 < segments.len() {
            name.push_str(&segment.to_string());
            continue;
        }
        let leaf = match segment {
            Segment::Name("weights") | Segment::Name("scale") => "weight",
            Segment::Name("shift") | Segment::Name("bias") => "bias",
            other => panic!("no checkpoint spelling for the leaf `{other}`"),
        };
        name.push_str(leaf);
    }
    name
}

/// Returns the state carrying the checkpoint: every parameter of
/// `model`'s tree restored by name, converted into the tree's element
/// type at the precision boundary.
///
/// This is the generic loader the module tier promises: the model
/// announces its own paths, [`foreign_name`] renders each as the
/// checkpoint's spelling, and [`named_restore`] does the rest —
/// missing tensors and shape mismatches fail loudly through the
/// restore's existing validation.
pub fn load<E: Element + From<f32>>(
    parameters: &Parameters<E>,
    model: &Gpt2<E>,
    weights: &Weights,
) -> Parameters<E> {
    let entries: Vec<(Path, Tensor<E>)> = named_parameters(model)
        .into_iter()
        .map(|(path, _)| {
            let payload = weights.tensor(&foreign_name(&path)).convert::<E>();
            (path, payload)
        })
        .collect();
    named_restore(parameters, model, entries)
}