cogforge-engine 2.1.0

A custom autograd engine and Transformer block built from scratch.

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. The goal is still to understand every gradient that flows — speed is a bonus, not the point.

---

Table of contents

• 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
• Gotchas
• Roadmap
• License

---

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

The package is organized into three 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.

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.

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.

---

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 broadcasting support.
a @ b	Batched matmul; gradients are correctly un-broadcast.
a[key]	Indexing/slicing.
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.
.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. Mind the distinction between losses that take probabilities and losses that take raw logits — this is the most common mistake.

Loss	Input expectation	Use when
Tensor.softmax_cross_entropy(scores, targets)	scores are raw logits, targets one-hot. Softmax fused inside (stable). Works for 2-D (B,V) and 3-D (B,T,V).	Standard classification / LM. Recommended.
Tensor.sparse_softmax_cross_entropy(scores, target_ids)	scores raw logits (B,T,V), target_ids integers (B,T).	Language modeling — skips building one-hot targets.
Tensor.sparse_softmax_cross_entropy_index(scores, labels, mask)	Raw logits + integer labels + per-token mask (1 = real, 0 = pad).	Padded LM / seq2seq batches. Recommended for Seq2Seq training.
Tensor.softmax_cross_entropy_masked(scores, targets, mask)	Logits + one-hot targets + per-row mask.	Padded batches, fused softmax.
Tensor.cross_entropy_loss(predictions, targets)	predictions are probabilities (call .softmax() first), targets one-hot.	You already have a softmax in your graph.
Tensor.cross_entropy_loss_masked(predictions, targets, mask)	Probabilities + per-row mask.	Padded batches without fused softmax.

ℹ️ 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.

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>; stops early when every sequence in the batch emits eos_id.
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(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
out = model.generate(src, sos_id=SOS, eos_id=EOS, max_new=32,
                     temperature=1.0, top_k=5)

---

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.

---

Roadmap

Shipped since the last release:

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

• 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)
• Checkpoint save/load for the transformer models

---

License

MIT License. See LICENSE for details.