Skip to main content

cogforge

A from-scratch deep learning library built on NumPy — a reverse-mode autograd engine extended all the way to working GPTs and encoder–decoder transformers, with optional GPU acceleration.

cogforge is a small, readable, educational deep learning framework. At its core is a Tensor that records every operation into a computation graph and backpropagates through it (micrograd-style), but unlike a toy autograd it scales up to real architectures: MLPs, RNNs, batch/layer normalization, multi-head self- and cross-attention, rotary position embeddings (RoPE), decoder-only GPTs, and a full encoder–decoder Seq2Seq transformer you can actually train and sample from.

There is no C++, no PyTorch — just NumPy and explicit, hand-derived gradients. Optionally, the entire backend can be swapped to CuPy for GPU execution, or accelerated with numexpr on CPU, without changing any model code. Every tensor also records the operation that produced it, and the built-in Graphviz visualizer can render the full computation graph — data, gradients, and ops — so you can literally see backpropagation. The goal is still to understand every gradient that flows — speed is a bonus, not the point.

This release overhauls the loss API. Seven overlapping cross-entropy variants have been unified into three functions covering every use case — with built-in label smoothing, optional padding masks everywhere, and one consistent shape/normalization convention. The superseded losses live on in cogforge.legacy. See What's new.


Table of contents


What's new

This version consolidates the loss API into three functions with a single, consistent convention — and adds label smoothing on the road to a full translator.

Unified cross-entropy API (3 losses, one convention)

Seven overlapping cross-entropy variants (softmax_cross_entropy_old, softmax_cross_entropy_masked, sparse_softmax_cross_entropy, sparse_softmax_cross_entropy_legacy, sparse_softmax_cross_entropy_index_legacy, cross_entropy_loss, cross_entropy_loss_masked) are replaced by three that tile the whole space with no overlap — see Losses:

  • Tensor.sparse_softmax_cross_entropy_index(scores, labels, mask=None, eps=0.0) — the workhorse: logits in, integer labels, optional padding mask, optional label smoothing via eps. No one-hot matrix is ever built.
  • Tensor.softmax_cross_entropy(scores, targets, mask=None) — logits in, arbitrary target distributions (knowledge distillation, mixup, soft labels), optional mask.
  • Tensor.cross_entropy_from_probs(predictions, targets, mask=None) — for graphs that already end in .softmax(); prefer the logits-based losses when you have logits.

All three now share the same rules, closing several long-standing traps:

  • One shape convention. Every loss accepts (N, V) and (B, T, V) logits (auto-flattened internally); masks may be (N,) or (B, T).
  • One normalization convention. Every loss divides by the number of real (unmasked) tokens — no more /B vs /B·T vs /n_real drift between variants, which silently changed effective learning rates when switching losses.
  • Masks are optional everywhere, defaulting to "all positions are real".
  • Exact log-probabilities. Loss values are computed from log-softmax directly instead of log(clip(p)), so near-zero probabilities are no longer floored at log(1e-15).

Label smoothing

sparse_softmax_cross_entropy_index(..., eps=0.1) trains against (1−eps) on the true class and eps/(V−1) spread uniformly over the rest — the standard regularizer for translation models — at zero extra memory: both loss and gradient are computed straight from the label indices. eps=0.0 (default) is exact plain cross-entropy.

Performance

  • The hot path (eps=0, no mask — e.g. char-level GPT training) keeps the fused, numexpr-accelerated in-place backward from the previous sparse_softmax_cross_entropy.
  • Masked/smoothed paths fold the mask, upstream gradient, and normalization into a single per-row coefficient — one fewer (N, V) temporary in every backward.

Legacy module

The superseded losses were moved, not deleted: they live in cogforge.legacy under TensorLegacy, unchanged, for reference and for reproducing old experiments. New code should not import them — note that they retain their old (inconsistent) normalization conventions.


Previous release — autograd correctness pass

This version focused on autograd correctness and debuggability. All fixes below were validated with finite-difference gradient checks and end-to-end training tests.

Gradient correctness fixes

  • Broadcast gradients fully generalized. +, -, and * previously mishandled un-broadcasting when shapes broadcast along an interior size-1 dimension (e.g. (B, 1, C) + (B, T, C) — exactly the shape of a padding mask or any keepdims tensor): depending on shapes this either crashed or silently accumulated wrong gradients. All elementwise ops now route through a single, verified Tensor.unbroadcast, matching @.
  • Indexing with repeated indices now accumulates. tensor[key]'s backward used direct fancy-index assignment, which drops gradient contributions for duplicate indices due to NumPy write buffering. It now uses the backend-aware scatter_add (np.add.at on CPU, cupyx.scatter_add on GPU), so gathers with repeated indices — RoPE slicing, token lookups — backprop correctly.
  • sparse_softmax_cross_entropy backward rewritten. The gradient is now written straight into scores.grad as (softmax − onehot) / N without copying or mutating the cached softmax output — faster, less memory, and verified against numerical gradients (max error ~1e-4).
  • BatchNorm1D is now safe to call multiple times before backwards(). Backward state (x_hat, std_inv, the training flag) was previously stored on the module, so a second forward pass clobbered what the first pass's backward needed. Each call's backward closure now captures its own state — interleaved forwards produce bit-identical gradients to isolated ones.
  • NO_GRAD fast paths are consistent everywhere: no op builds a graph node it's about to discard.

Generation fixes

  • Seq2Seq.generate handles per-sequence EOS. Previously the loop only stopped if every row in the batch emitted eos_id on the same step, and finished rows kept sampling. Each row is now frozen to pad_id the moment it emits EOS, and the loop exits once all rows are done.

New: graph visualization (cogforge.utils)

  • Every Tensor now carries a human-readable .op label ("+", "MATMUL", "softmax", "layernorm", …) recording the operation that produced it.
  • draw_graph(tensor) renders the full computation graph via Graphviz — one record node per tensor showing its op, shape, data, and gradient. See Visualizing the computation graph.

Installation

pip install cogforge-engine

Requires Python 3.8+ and NumPy — that's the only hard dependency. Two optional extras unlock acceleration:

pip install cupy-cuda12x   # GPU backend (pick the build matching your CUDA version)
pip install numexpr        # multi-threaded CPU element-wise ops
pip install graphviz       # computation-graph visualization (also needs the Graphviz system binaries)

The package is organized into four modules:

Module Contains
cogforge.backend The swappable array backend: NumPy ↔ CuPy switching, numexpr flag, global no-grad flag.
cogforge.app The autograd engine (Tensor) and every building block — layers, optimizers, losses, normalization, attention, positional encodings.
cogforge.models Ready-to-use models: GPTV1, GPT2, Seq2Seq.
cogforge.utils Computation-graph tracing and Graphviz rendering (trace_graph, draw_graph).
cogforge.legacy TensorLegacy — the superseded cross-entropy variants, kept unchanged for reference and reproducing old experiments. Don't use in new code.
from cogforge.app import Tensor, Linear, Adam, MultiHeadAttention   # building blocks
from cogforge.models import GPTV1, GPT2, Seq2Seq                    # models
from cogforge import backend                                        # device control

Quick start

import numpy as np
from cogforge.app import Tensor

# Build a graph
a = Tensor(np.array([2.0, 3.0]))
b = Tensor(np.array([4.0, 5.0]))
c = (a * b).sigmoid().softmax()

# Backpropagate (note the spelling: backwards, with an 's')
c.backwards()

print(a.grad)   # gradient of the output w.r.t. a

Every Tensor carries a .data (the backend array), a .grad (same shape, accumulates gradients), and a hidden _backwards closure that knows how to push gradient to its parents. Calling .backwards() on any node runs a topological sort and walks the graph in reverse.


Backend: CPU, GPU, numexpr, and no-grad mode

cogforge.backend exposes a module-level np that every layer and model routes through. By default it is NumPy; flipping one switch reroutes the whole library to CuPy.

GPU (CuPy)

from cogforge import backend
backend.use_gpu(True)     # everything created after this lives on the GPU
# ... build model, train ...
backend.use_gpu(False)    # back to NumPy
  • Raises RuntimeError if CuPy is not installed.
  • Switch before constructing your model — parameters are allocated on whichever device is active at creation time.
  • Embedding's scatter-add backward automatically uses cupyx.scatter_add on GPU and np.add.at on CPU.
  • Sampling in generate() always happens on CPU (logits are pulled back with to_cpu), so generation works identically on either device.

numexpr (CPU acceleration)

from cogforge.app import set_numexpr
set_numexpr(True, threads=8)   # multi-threaded softmax / fused elementwise ops
set_numexpr(False)

Only takes effect on the CPU backend (ignored when the GPU is active). Raises if numexpr isn't installed.

No-grad mode

from cogforge.app import needGradientHence
needGradientHence(False)   # stop building graphs: no .grad buffers, no closures
# ... fast inference ...
needGradientHence(True)    # back to training mode

When gradients are off, every op returns a bare result tensor — no children, no backward closure, no gradient buffers — which slashes memory use and speeds up inference. All three models' generate() methods toggle this automatically and restore the previous state afterwards (in a try/finally, so it's restored even on error).

Helpers

Function Purpose
to_cpu(a) Return a NumPy array regardless of the active backend. Use it before plotting, sampling, or saving.
scatter_add(target, indices, values) Backend-aware target[indices] += values (handles repeated indices correctly on both devices).

Core concept: the Tensor

Tensor(array, children=(), requires_grad=True, typed="compressed")
Argument Meaning
array Any array-like; stored on the active backend in .data.
children Parent tensors in the graph (set internally by ops; you rarely pass this).
requires_grad Reserved flag (currently informational).
typed "compressed" → float32 (default), anything else → float64.
op Human-readable label of the operation that produced this tensor ("+", "MATMUL", "softmax", …). Set internally by every op; used by the graph visualizer and handy when debugging.

Gradients accumulate into .grad. Always zero them between optimization steps (the optimizers do this for you via zero_grad()). When global no-grad mode is on, .grad is None and no graph is recorded.

The graph is single-use. backwards() frees the graph as it goes (clearing each node's children and backward closure) to release memory eagerly. Run one forward → one backward per step, and if you want to visualize the graph, draw it before calling backwards().


API reference

Tensor — autograd engine

Differentiable operations (each builds graph and defines its own backward):

Operation Notes
a + b, a - b, a * b Elementwise, with full broadcasting support — gradients are correctly un-broadcast for leading and interior size-1 dimensions.
a @ b Batched matmul; gradients are correctly un-broadcast.
a[key] Indexing/slicing. Backward uses backend-aware scatter-add, so repeated indices accumulate gradients correctly on CPU and GPU.
Tensor.cat(tensors, axis=-1) Classmethod. Concatenates a tuple of tensors along axis; backward splits the gradient back to each parent. (Used internally by RoPE.)
.relu()
.sigmoid()
.tanh()
.softmax(axis=-1) Numerically stable (max-subtraction); numexpr-accelerated when enabled.
.view(shape) Reshape (handles non-contiguous data).
.flatten() Flattens everything after the batch dim → (B, -1).
.flatten_consective(num) Groups num consecutive timesteps. Expects a 3-D (B, T, C) tensor; T must be divisible by num.
.transpose(axes) Permute axes (pass the full permutation tuple).
.masked_fill(mask, value) Sets entries where mask is True to value (used for causal/padding attention masks).
.dropout(p=0.1, training=True) Inverted dropout: scales by 1/(1-p) at train time, identity when training=False.
.dropTheWholeNeuron(p=0.1, training=True, axis=-1, batch_ind=0) Structured dropout — zeroes entire feature channels rather than individual elements.

Backward pass

Method Notes
.backwards() Primary. Iterative topological sort — safe for deep/long graphs. Frees the graph as it runs (single-use; see the note under Core concept).
.backwards_recursive() Legacy recursive version; can hit Python's recursion limit on long sequences. Prefer .backwards().

Static helper

  • Tensor.unbroadcast(grad, shape) — reduces a broadcasted gradient back to the original parameter shape. Used internally.

Losses

All losses are classmethods on Tensor and return a scalar loss tensor you call .backwards() on. There are exactly three, and they share one convention:

  • Shapes: logits/probs may be (N, V) or (B, T, V) — 3-D input is auto-flattened. Masks may be (N,) or (B, T).
  • Masks: optional everywhere (mask=None treats all positions as real); 1 = real token, 0 = padding. Padded positions contribute zero loss and exactly zero gradient.
  • Normalization: always by the number of real tokens (all positions when unmasked).

Pick by what your target looks like and what your input is:

Loss Input Target Use when
Tensor.sparse_softmax_cross_entropy_index(scores, labels, mask=None, eps=0.0) raw logits integer class ids The workhorse. Classification, LM, seq2seq. eps>0 enables label smoothing: 1−eps on the true class, eps/(V−1) on the rest — computed straight from indices, no one-hot built. eps=0.1 is the standard for translation.
Tensor.softmax_cross_entropy(scores, targets, mask=None) raw logits full distribution (…, V), rows sum to 1 Targets that aren't expressible as an index: knowledge distillation against a teacher's output, mixup, annotator vote distributions.
Tensor.cross_entropy_from_probs(predictions, targets, mask=None) probabilities (e.g. output of .softmax()) one-hot or distribution Your graph already ends in a softmax. Prefer the logits-based losses when you have logits — their gradients are exact and bounded, this one saturates via clipping as probs → 0.

ℹ️ The first two apply softmax internally — feed them raw logits. cross_entropy_from_probs is the opposite — it expects probabilities. Mixing these up silently trains the wrong thing.

The previous generation of losses (softmax_cross_entropy_masked, sparse_softmax_cross_entropy, cross_entropy_loss, cross_entropy_loss_masked, and the _old/_legacy variants) now lives in cogforge.legacy under TensorLegacy, unchanged. They keep their old, mutually inconsistent normalizations (/B vs /B·T vs /n_real) — use them only to reproduce old experiments, and don't mix them with the current losses in one training run.


Layers

Linear(nin, nout)

Affine transform x @ W + b. He-initialized weights. .parameters() → [W, b].

Embedding(vocab_size, embedding_dim)

Lookup table. Call with an integer index array; backward scatters gradients correctly (repeated indices accumulate, on CPU and GPU). .parameters() → [weights].

LayerNorm(dim, eps=1e-5)

Normalizes over the last dimension. Learnable gamma/beta, full hand-derived backward. .parameters() → [gamma, beta].

BatchNorm1D(dim, eps=1e-5, momentum=0.1)

Normalizes over the batch (and time, for 3-D input). Tracks running_mean/running_var for inference. Toggle .training = True/False. Learnable gamma/beta. Each forward call captures its own backward state, so calling the layer multiple times before backwards() (e.g. gradient accumulation, shared modules) yields correct gradients for every call.

FeedForward(dmodel, dff=None)

Position-wise MLP: Linear → ReLU → Dropout(0.15) → Linear. dff defaults to 4 * dmodel. Dropout is active only when called with is_training=True (transformer blocks handle this for you via their train/infer state).


Attention

Attention(dk)

Scaled dot-product attention. Call attention(Q, K, V, mask=None). dk sets the 1/√dk scale. Masked positions are filled with -1e9 before the softmax.

MultiHeadAttention(dinp, dmodel, dout, n, rope=None)

n heads, dmodel split into n chunks of size dmodel // n (must divide evenly). Projects input dinp → dmodel, attends, projects dmodel → dout. If a rope (see RotatoryPositionalEncoding) is passed, it is applied to Q and K after the head split — this is how GPT2 gets rotary positions. Call mha(query, key, value, mask=None). .parameters() returns all four projection layers' params.

CrossAttention(dim_dec, dim_enc, d_out, dec_rope=None, enc_rope=None, d_k=None, h=None, d_model=None)

Attention where queries come from the decoder stream and keys/values from the encoder stream — the bridge of an encoder–decoder transformer. Specify head geometry as either (d_k and h) or (d_model and h). Optional separate RoPE for the query (decoder) side and key (encoder) side. Call cross(x_decod, x_encod, mask=None); pass the encoder padding mask as mask so the decoder never attends to pad tokens.


Positional encodings

PositionalEncoding(max_len, dmodel)

Fixed sinusoidal positions, added to the input embeddings. Call pe(x). No parameters. Used by GPTV1.

RotatoryPositionalEncoding(max_len, dim, base=10000.0)

Rotary position embeddings (RoPE). Instead of adding position vectors to embeddings, it rotates Q and K inside attention, encoding relative position directly in the dot product. dim is the per-head dimension d_k (must be even), not d_model. Construct once and hand the same instance to every block:

rope = RotatoryPositionalEncoding(max_len, d_model // n_heads)
block = Transformer(dmodel=d_model, n=n_heads, rope=rope)

No parameters. Used by GPT2 and (optionally) Seq2Seq.


Transformer blocks

Transformer(dmodel, n, dff=None, rope=None, is_training=False)

A pre-norm self-attention block: x + Attn(LN(x)) then x + FF(LN(x)). n = number of heads. Optional RoPE. Call block(x, mask=None) — pass a causal mask for LM use or a padding mask for encoder use.

State control: .train(enabled=True) / .infer(enabled=True) toggle is_training, which switches the feed-forward dropout on/off.

Decoder(d_model, n_heads, d_ff=None, is_training=False, dec_in_rope=None, enc_rope=None, dec_rope=None)

A full pre-norm encoder–decoder block with three sublayers:

  1. masked self-attention over the decoder stream (dec_in_rope optional),
  2. cross-attention into the encoder output (dec_rope on queries, enc_rope on keys, both optional),
  3. feed-forward with dropout.

Call block(x_dec, x_enc, mask=None, cross_mask=None) — mask is the causal mask for self-attention, cross_mask the encoder padding mask. Same .train() / .infer() interface as Transformer.


Containers

Sequential(layers)

Runs layers in order. .train() / .test() flip the training flag on any layer that has one (e.g. BatchNorm1D).

⚠️ Sequential.parameters() only collects layers exposing W, b, gamma, or beta attributes (i.e. Linear, LayerNorm, BatchNorm1D). Composite layers like MultiHeadAttention, FeedForward, and Transformer hold sub-modules, so their parameters are not picked up here — gather those via each module's own .parameters().

MLP(layer_sizes)

Convenience feed-forward net: Linear → ReLU between layers, plain Linear output. Built from a list of sizes, e.g. MLP([784, 128, 64, 10]).

  • .save(filename="best_model.npz") / .load(filename="best_model.npz") — persist/restore weights.
  • Note: MLP does not expose a parameters() method; collect them via [p for layer in mlp.layers for p in layer.parameters()] if you want to optimize it.

Recurrent

RNNCell(input_dim, hidden_dim)

One tanh recurrence step: h_next = tanh(i2h(x) + h2h(h_prev)). .parameters() included.

RNN(input_dim, hidden_dim)

Unrolls a cell over a list of timestep tensors (each (B, input_dim)) and returns the list of hidden states (each (B, hidden_dim)). Optional prev_hidden.

StackedRNN(input_dim, hidden_dim, num_layers)

Multiple RNN layers stacked. Returns (top_layer_states, per_layer_final_states) — the second value is convenient for seq2seq.

Bridge(enc_hidden, dec_hidden, enc_layers, dec_layers, mode="project")

Maps RNN encoder final hidden states to decoder initial hidden states, handling mismatched layer counts and hidden sizes.

mode Behavior
"project" One learned Linear(enc_hidden → dec_hidden) per decoder layer. General, recommended.
"tie" No parameters; requires enc_hidden == dec_hidden. Selects/repeats raw states.

Optimizers

Both take an iterable of parameter tensors and share the same interface: step(), zero_grad(), clip_grads(max_norm=5.0).

SGD(parameters, learning_rate=0.01)

Plain stochastic gradient descent.

Adam(parameters, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8)

Adam with bias correction. Recommended for transformers.

opt = Adam(model.parameters(), lr=3e-4)
opt.zero_grad()
loss.backwards()
opt.clip_grads(1.0)   # optional gradient clipping
opt.step()

Models

All three models share conventions:

  • model.parameters() returns every trainable tensor (deduplicated where weights are shared).
  • model.generate(...) automatically switches to inference mode and disables gradient tracking for the duration, restoring the previous state afterwards — you never need to toggle anything manually to sample.
  • Sampling supports temperature and top_k, is numerically stabilized, and always runs on CPU regardless of backend.

GPTV1(vocab, d_model, n_heads, n_layers, max_len, d_ff=None)

A decoder-only transformer with sinusoidal (additive) positional encoding: token embedding + positions + stacked pre-norm Transformer blocks + final LayerNorm + output head. Causal masking is applied internally.

Method Description
model(idx) idx: integer array (B, T). Returns logits (B, T, vocab).
model.generate(idx, n_new, temperature=1.0, top_k=None) Autoregressive sampling. Crops the context to the last max_len tokens. Returns (B, T + n_new).

GPT2(vocab, d_model, n_heads, n_layers, max_len, d_ff=None, base=10000.0, training=False)

The modern decoder-only variant: RoPE instead of additive positions (one shared RotatoryPositionalEncoding of dim d_model // n_heads applied to Q/K in every block), no positional add at the input, dropout in the feed-forward layers when training. base is the RoPE frequency base.

Method Description
model(idx) Logits (B, T, vocab) with causal masking applied internally.
model.train() / model.infer() Toggle training mode (dropout on/off) across all blocks.
model.generate(idx, n_new, temperature=1.0, top_k=None) As GPTV1; also handles the train/infer switch for you.

Seq2Seq(enc_vocab, dec_vocab, d_model, n_heads, num_enc_layers, num_dec_layers, max_len, d_ff=None, training=False, shared_tok=False, pad_id=0, encoder_rope=None, dec_in_rope=None, dec_rope=None, enc_rope=None)

A full encoder–decoder transformer (the original Attention Is All You Need topology, pre-norm):

  • Encoder: num_enc_layers self-attention Transformer blocks over the source, with a padding mask built from pad_id, followed by a final encoder LayerNorm.
  • Decoder: num_dec_layers Decoder blocks — causal self-attention, cross-attention into the encoder output (respecting the encoder padding mask), feed-forward.
  • Weight tying: with shared_tok=True and enc_vocab == dec_vocab, the encoder embedding, decoder embedding, and output projection all share one matrix (embeddings scaled by √d_model, plus a learned output bias). Cuts parameter count substantially.
  • RoPE, opt-in per site: pass any non-None value to encoder_rope (encoder self-attention), dec_in_rope (decoder self-attention), dec_rope (cross-attention queries), and/or enc_rope (cross-attention keys) to enable a shared rotary encoding at that site.
Method Description
model(enc_idx, dec_idx) Teacher-forced forward. enc_idx: source (B, T_enc); dec_idx: shifted target starting with <SOS>, (B, T_dec). Returns logits (B, T_dec, dec_vocab).
model.encode(enc_idx) Run the encoder once; returns (x_enc, enc_pad_mask) for reuse across decode steps.
model.decode_step(dec_idx, x_enc, enc_pad) Decoder forward against a fixed encoder output.
model.generate(enc_idx, sos_id, eos_id=None, max_new=50, temperature=1.0, top_k=None) Encodes once, then autoregressively decodes from <SOS>. Tracks completion per sequence: the moment a row emits eos_id it is frozen and padded with pad_id for the remaining steps, and decoding stops early once every row has finished.
model.train() / model.infer() Toggle dropout across all encoder and decoder blocks.
Seq2Seq.make_pad_mask(idx, pad_id) Static helper: (B, T) ints → (B, 1, 1, T) boolean mask, True at padding.

Worked example: train a char-level GPT

Works identically with GPTV1; shown with the RoPE-based GPT2.

import numpy as np
from cogforge.app import Tensor, Adam
from cogforge.models import GPT2

# --- data -------------------------------------------------------------
text  = open("input.txt").read()
chars = sorted(set(text))
stoi  = {c: i for i, c in enumerate(chars)}
itos  = {i: c for i, c in enumerate(chars)}
data  = np.array([stoi[c] for c in text])
vocab = len(chars)

# --- model ------------------------------------------------------------
block = 64
model = GPT2(vocab=vocab, d_model=128, n_heads=4,
             n_layers=4, max_len=block, training=True)
opt   = Adam(model.parameters(), lr=3e-4)

def get_batch(bs=32):
    ix = np.random.randint(0, len(data) - block - 1, size=bs)
    x  = np.stack([data[i:i + block]     for i in ix])
    y  = np.stack([data[i + 1:i + block + 1] for i in ix])
    return x, y

# --- train ------------------------------------------------------------
for step in range(2000):
    x, y   = get_batch()
    logits = model(x)                                  # (B, T, vocab)
    loss   = Tensor.sparse_softmax_cross_entropy_index(logits, y)

    opt.zero_grad()
    loss.backwards()
    opt.clip_grads(1.0)
    opt.step()

    if step % 100 == 0:
        print(f"step {step:4d} | loss {float(loss.data):.4f}")

# --- sample -----------------------------------------------------------
ctx = np.array([[stoi["\n"]]])
out = model.generate(ctx, n_new=300, temperature=0.8, top_k=20)
print("".join(itos[int(i)] for i in out[0]))

To run the same script on GPU, add two lines at the top — before building the model:

from cogforge import backend
backend.use_gpu(True)

Worked example: encoder–decoder Seq2Seq

A toy translation/unscrambling setup with a shared vocabulary, tied weights, and RoPE everywhere:

import numpy as np
from cogforge.app import Tensor, Adam
from cogforge.models import Seq2Seq

PAD, SOS, EOS = 0, 1, 2
vocab = 40

model = Seq2Seq(
    enc_vocab=vocab, dec_vocab=vocab,
    d_model=128, n_heads=4,
    num_enc_layers=3, num_dec_layers=3,
    max_len=32, training=True,
    shared_tok=True, pad_id=PAD,
    encoder_rope=True, dec_in_rope=True,   # any non-None value enables RoPE at that site
)
opt = Adam(model.parameters(), lr=3e-4)

for step in range(num_steps):
    src, tgt = get_batch()                 # src: (B, T_enc) padded with PAD
                                           # tgt: (B, T_dec+1) = [SOS, ..., EOS, PAD...]
    dec_in, labels = tgt[:, :-1], tgt[:, 1:]
    logits = model(src, dec_in)            # (B, T_dec, vocab)

    mask = (labels != PAD).astype(np.float32)     # (B, T_dec): 1 = real, 0 = pad
    loss = Tensor.sparse_softmax_cross_entropy_index(logits, labels,
                                                     mask=mask, eps=0.1)  # label smoothing

    opt.zero_grad()
    loss.backwards()
    opt.clip_grads(1.0)
    opt.step()

# inference: encode once, decode token by token, stop on EOS (per sequence)
out = model.generate(src, sos_id=SOS, eos_id=EOS, max_new=32,
                     temperature=1.0, top_k=5)

Visualizing the computation graph

cogforge.utils renders the autograd graph with Graphviz — every node shows the op that produced it, its shape, its data, and its gradient. It's the fastest way to see what your model is doing and to debug shape or gradient-flow issues.

import numpy as np
from cogforge.app import Tensor, Linear
from cogforge.utils import draw_graph

x   = Tensor(np.random.randn(2, 4))
lin = Linear(4, 3)
out = lin(x).relu().softmax()

dot = draw_graph(out, rankdir='LR')   # draw BEFORE backwards() — the graph is freed by it
dot.render('graph', view=True)        # writes graph.svg and opens it
Function Purpose
trace_graph(root_tensor) Walks the graph from a root tensor; returns (nodes, edges).
draw_graph(root_tensor, max_char=50, format='svg', rankdir='LR') Builds a Graphviz Digraph. max_char truncates long data/grad strings; format is any Graphviz output format (svg, png, pdf); rankdir='TB' for top-to-bottom layout.

Nodes produced by a named op render white with the op label; raw leaf tensors (parameters, inputs) render light blue. Requires the graphviz Python package and the Graphviz system binaries (apt install graphviz / brew install graphviz).

Two practical notes: call draw_graph before backwards(), since the backward pass frees the graph as it runs (afterwards you'll see a single orphan node); and keep the visualized graph small — a full transformer forward produces hundreds of nodes, so draw a single block or a toy input rather than a whole training step.


Gotchas

  • It's backwards(), not backward(). The backward pass method has a trailing s.
  • Logits vs. probabilities. softmax_cross_entropy and sparse_softmax_cross_entropy_index fuse the softmax internally — feed them raw logits. cross_entropy_from_probs expects probabilities. Mixing these up silently trains the wrong thing.
  • Label smoothing raises the loss floor. With eps > 0 the minimum achievable loss is no longer 0 (a perfect model still pays the smoothing term), so judge training by token accuracy or validation metrics, not by how close the raw loss gets to zero — and don't compare loss values across different eps settings.
  • Legacy losses normalize differently. cogforge.legacy.TensorLegacy keeps the old /B / /B·T conventions; the current losses always divide by real-token count. Swapping one for the other rescales gradients — retune the learning rate if you migrate an old script.
  • Gradients accumulate. Call optimizer.zero_grad() every step (or p.grad[...] = 0), or gradients pile up across iterations.
  • Switch the backend before building the model. use_gpu(True) after construction leaves your parameters stranded on the CPU while new activations land on the GPU.
  • RoPE dim is per-head. RotatoryPositionalEncoding takes d_model // n_heads (which must be even), not d_model. The models handle this internally — it only matters if you wire blocks up yourself.
  • Training vs. inference mode matters now. GPT2 and Seq2Seq use dropout; call .train() before optimizing and .infer() before evaluating. generate() handles this (and no-grad mode) for you and restores the prior state afterwards.
  • Sequential.parameters() is shallow — see the note under Containers. For attention/feed-forward/transformer stacks, gather parameters through each module's own .parameters() (as the models' parameters() methods do).
  • RNNs operate on lists, not a single (B, T, C) tensor — pass a list of per-timestep tensors.
  • Don't dedupe tied parameters yourself. Seq2Seq.parameters() already deduplicates shared tensors by identity, so the tied embedding is only updated once per step.
  • The graph is single-use. backwards() frees children and backward closures as it runs. One forward → one backward. Calling backwards() twice on the same graph is a silent no-op the second time, and draw_graph must be called before, not after.
  • Visualize small graphs. draw_graph on a full model forward will produce an unreadable diagram with hundreds of nodes; visualize a single layer or block instead.

Roadmap

Shipped in this release:

  • ✅ Unified loss API — three cross-entropy losses covering index targets, distribution targets, and probability inputs, with one shape and normalization convention (see What's new)
  • ✅ Label smoothing (eps) in sparse_softmax_cross_entropy_index, computed from indices with zero extra memory
  • ✅ Optional padding masks and (B, T, V) auto-flattening in every loss
  • ✅ cogforge.legacy module (TensorLegacy) preserving the superseded losses for reproducibility

Shipped previously:

  • ✅ Autograd correctness pass — all gradients verified against finite-difference checks; broadcast, indexing, sparse cross-entropy, and BatchNorm backward bugs fixed
  • ✅ Computation-graph visualization via Graphviz (cogforge.utils), with per-tensor op labels
  • ✅ Per-sequence EOS handling in Seq2Seq.generate (finished rows freeze to pad_id)
  • ✅ Eager graph freeing in backwards() for lower peak memory
  • ✅ RoPE (rotary position embeddings), usable in GPT and at every attention site of the Seq2Seq model
  • ✅ Full encoder–decoder transformer (Seq2Seq) with cross-attention and padding masks
  • ✅ Weight tying between embeddings and the output head
  • ✅ GPU backend via CuPy; numexpr-accelerated CPU ops
  • ✅ Dropout (element-wise and structured) with train/infer modes
  • ✅ Global no-grad mode for fast, memory-light inference

Planned / under consideration:

  • Beam search decoding for Seq2Seq (length-normalized, per-sequence finished pool)
  • Lazy Gradient Allocation (LGA) for lower peak memory on long sequences
  • Cosine LR function as part of UTILS module
  • Weight decay / AdamW optimizer
  • Checkpoint save/load (state_dict / load_state_dict) for the transformer models, including optimizer state
  • no_grad() context manager and a shared Module base class
  • KV cache for faster generation
  • SwiGLU feed-forward and RMSNorm
  • RoPE length interpolation
  • Linear-attention block (as a study in the recall-vs-cost tradeoff)

License

APACHE License. See LICENSE for details.

Release files for cogforge-engine 2.1.3

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for cogforge-engine 2.1.3
File Size Uploaded
cogforge_engine-2.1.3.tar.gz 44.5 kB Details

Built distribution (wheel)

Table of built distributions (wheels) for cogforge-engine 2.1.3
File Interpreter ABI Platform
cogforge_engine-2.1.3-py3-none-any.whl Python 3 none any Details

Total release size: 78.7 kB

Release files / cogforge_engine-2.1.3.tar.gz

Download URL cogforge_engine-2.1.3.tar.gz
Size 44.5 kB
Tags Source
SHA-256 checksum
How to use checksums
c8c0f50587f58104b383bd7beddf444ea51e965bb0a9dd123c690643622bfb9a
BLAKE2b-256 checksum
How to use checksums
377426196e063b2f6e92df81aa45daf811678d496627d7c6cc3524fe7879d2b0
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
Yes
Uploaded via twine/6.1.0 CPython/3.13.12

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

PyPI Publish Attestation

PyPI verified that this artifact, at this checksum, originated from the publisher listed below.

Signed by GitHub Actions, verified by PyPI on Jul 7, 2026.

Transparency log

Release files / cogforge_engine-2.1.3-py3-none-any.whl

Download URL cogforge_engine-2.1.3-py3-none-any.whl
Size 34.2 kB
Tags Python 3
SHA-256 checksum
How to use checksums
53b4e1d437b3faa377a353235f2bc878b2d644d667fd3ba5ec2fff294d3487e8
BLAKE2b-256 checksum
How to use checksums
d82ee3c4e69d207ebd63017358c6230eddead34ca7426e6899680fe97fe74a65
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
Yes
Uploaded via twine/6.1.0 CPython/3.13.12

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

PyPI Publish Attestation

PyPI verified that this artifact, at this checksum, originated from the publisher listed below.

Signed by GitHub Actions, verified by PyPI on Jul 7, 2026.

Transparency log

Release history Release notifications | RSS feed

This release

2.1.3 This release

2 release files

2.1.2

2 release files

2.1.1

2 release files

2.1.0

2 release files

2.0.1

2 release files

2.0.0

2 release files

1.1.0

2 release files

1.0.7

2 release files

1.0.6

2 release files

1.0.5

2 release files

1.0.4

2 release files

1.0.3

2 release files

1.0.2

2 release files

1.0.1

2 release files

1.0.0

2 release files

0.2.1

2 release files

0.2.0

2 release files

0.1.0

2 release files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page