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 is a correctness pass. Every gradient in the engine has now been verified against finite-difference numerical checks, and several subtle autograd bugs were found and fixed — see What's new.

Table of contents

• What's new
• Installation
• Quick start
• Backend: CPU, GPU, numexpr, and no-grad mode
• Core concept: the Tensor
• API reference
  • Tensor — autograd engine
  • Losses
  • Layers
  • Attention
  • Positional encodings
  • Transformer blocks
  • Containers
  • Recurrent
  • Optimizers
  • Models
• Worked example: train a char-level GPT
• Worked example: encoder–decoder Seq2Seq
• Visualizing the computation graph
• Gotchas
• Roadmap
• License

What's new

This version focuses 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:

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

Core concept: the Tensor

Tensor(array, children=(), requires_grad=True, typed="compressed")

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):

Backward pass

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. Mind the distinction between losses that take probabilities and losses that take raw logits — this is the most common mistake.

ℹ️ The softmax_* and sparse_softmax_* variants apply softmax internally — feed them raw logits. cross_entropy_loss* is the opposite — it expects probabilities.

sparse_softmax_cross_entropy_legacy and softmax_cross_entropy_old are kept for reference; prefer the current versions.

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.

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.

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.

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.

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(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)
    loss = Tensor.sparse_softmax_cross_entropy_index(logits, labels, mask)

    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

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 / sparse_softmax_cross_entropy* fuse the softmax internally — feed them raw logits. cross_entropy_loss* expects probabilities. Mixing these up silently trains the wrong thing.
• 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:

• ✅ Autograd correctness pass — all gradients verified against finite-difference checks; broadcast, indexing, sparse cross-entropy, and BatchNorm backward bugs fixed (see What's new)
• ✅ 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

Shipped previously:

• ✅ 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:

• 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.