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
- Installation
- Quick start
- Backend: CPU, GPU, numexpr, and no-grad mode
- Core concept: the
Tensor - API reference
- Worked example: train a char-level GPT
- Worked example: encoder–decoder Seq2Seq
- Visualizing the computation graph
- Gotchas
- Roadmap
- License
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 viaeps. 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
/Bvs/B·Tvs/n_realdrift 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-softmaxdirectly instead oflog(clip(p)), so near-zero probabilities are no longer floored atlog(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 previoussparse_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 anykeepdimstensor): depending on shapes this either crashed or silently accumulated wrong gradients. All elementwise ops now route through a single, verifiedTensor.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-awarescatter_add(np.add.aton CPU,cupyx.scatter_addon GPU), so gathers with repeated indices — RoPE slicing, token lookups — backprop correctly. sparse_softmax_cross_entropybackward rewritten. The gradient is now written straight intoscores.gradas(softmax − onehot) / Nwithout copying or mutating the cached softmax output — faster, less memory, and verified against numerical gradients (max error ~1e-4).BatchNorm1Dis now safe to call multiple times beforebackwards(). 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_GRADfast paths are consistent everywhere: no op builds a graph node it's about to discard.
Generation fixes
Seq2Seq.generatehandles per-sequence EOS. Previously the loop only stopped if every row in the batch emittedeos_idon the same step, and finished rows kept sampling. Each row is now frozen topad_idthe moment it emits EOS, and the loop exits once all rows are done.
New: graph visualization (cogforge.utils)
- Every
Tensornow carries a human-readable.oplabel ("+","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
RuntimeErrorif 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 usescupyx.scatter_addon GPU andnp.add.aton CPU.- Sampling in
generate()always happens on CPU (logits are pulled back withto_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 callingbackwards().
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=Nonetreats 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_probsis 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/_legacyvariants) now lives incogforge.legacyunderTensorLegacy, unchanged. They keep their old, mutually inconsistent normalizations (/Bvs/B·Tvs/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:
- masked self-attention over the decoder stream (
dec_in_ropeoptional), - cross-attention into the encoder output (
dec_ropeon queries,enc_ropeon keys, both optional), - 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 exposingW,b,gamma, orbetaattributes (i.e.Linear,LayerNorm,BatchNorm1D). Composite layers likeMultiHeadAttention,FeedForward, andTransformerhold 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:
MLPdoes not expose aparameters()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
temperatureandtop_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_layersself-attentionTransformerblocks over the source, with a padding mask built frompad_id, followed by a final encoderLayerNorm. - Decoder:
num_dec_layersDecoderblocks — causal self-attention, cross-attention into the encoder output (respecting the encoder padding mask), feed-forward. - Weight tying: with
shared_tok=Trueandenc_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-
Nonevalue toencoder_rope(encoder self-attention),dec_in_rope(decoder self-attention),dec_rope(cross-attention queries), and/orenc_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(), notbackward(). The backward pass method has a trailings. - Logits vs. probabilities.
softmax_cross_entropyandsparse_softmax_cross_entropy_indexfuse the softmax internally — feed them raw logits.cross_entropy_from_probsexpects probabilities. Mixing these up silently trains the wrong thing. - Label smoothing raises the loss floor. With
eps > 0the 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 differentepssettings. - Legacy losses normalize differently.
cogforge.legacy.TensorLegacykeeps the old/B//B·Tconventions; 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 (orp.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.
RotatoryPositionalEncodingtakesd_model // n_heads(which must be even), notd_model. The models handle this internally — it only matters if you wire blocks up yourself. - Training vs. inference mode matters now.
GPT2andSeq2Sequse 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. Callingbackwards()twice on the same graph is a silent no-op the second time, anddraw_graphmust be called before, not after. - Visualize small graphs.
draw_graphon 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) insparse_softmax_cross_entropy_index, computed from indices with zero extra memory - ✅ Optional padding masks and
(B, T, V)auto-flattening in every loss - ✅
cogforge.legacymodule (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-tensoroplabels - ✅ Per-sequence EOS handling in
Seq2Seq.generate(finished rows freeze topad_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 sharedModulebase 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)
| File | Size | Uploaded | |
|---|---|---|---|
| cogforge_engine-2.1.3.tar.gz | 44.5 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| 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 logRelease 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