cogforge-engine 2.1.0__tar.gz → 2.1.2__tar.gz

{cogforge_engine-2.1.0 → cogforge_engine-2.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cogforge-engine
-Version: 2.1.0
+Version: 2.1.2
 Summary: A custom autograd engine and Transformer block built from scratch.
 Project-URL: Homepage, https://github.com/avikmjd2/cogforge
 Author-email: Avik Majumder <avikmjd2@gmail.com>
@@ -19,12 +19,15 @@ Description-Content-Type: text/markdown
 
 `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.
+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](#whats-new).
 
 ---
 
 ## Table of contents
 
+- [What's new](#whats-new)
 - [Installation](#installation)
 - [Quick start](#quick-start)
 - [Backend: CPU, GPU, numexpr, and no-grad mode](#backend-cpu-gpu-numexpr-and-no-grad-mode)
@@ -42,12 +45,36 @@ There is no C++, no PyTorch — just NumPy and explicit, hand-derived gradients.
   - [Models](#models)
 - [Worked example: train a char-level GPT](#worked-example-train-a-char-level-gpt)
 - [Worked example: encoder–decoder Seq2Seq](#worked-example-encoderdecoder-seq2seq)
+- [Visualizing the computation graph](#visualizing-the-computation-graph)
 - [Gotchas](#gotchas)
 - [Roadmap](#roadmap)
 - [License](#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](#visualizing-the-computation-graph).
+
+---
+
 ## Installation
 
 ```bash
@@ -59,15 +86,17 @@ Requires Python 3.8+ and NumPy — that's the only hard dependency. Two optional
 ```bash
 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 three modules:
+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`). |
 
 ```python
 from cogforge.app import Tensor, Linear, Adam, MultiHeadAttention   # building blocks
@@ -158,9 +187,12 @@ Tensor(array, children=(), requires_grad=True, typed="compressed")
 | `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](#visualizing-the-computation-graph), draw it **before** calling `backwards()`.
+
 ---
 
 ## API reference
@@ -171,9 +203,9 @@ Gradients **accumulate** into `.grad`. Always zero them between optimization ste
 
 | Operation | Notes |
 | --- | --- |
-| `a + b`, `a - b`, `a * b` | Elementwise, with broadcasting support. |
+| `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. |
+| `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()` | |
@@ -191,7 +223,7 @@ Gradients **accumulate** into `.grad`. Always zero them between optimization ste
 
 | Method | Notes |
 | --- | --- |
-| `.backwards()` | **Primary.** Iterative topological sort — safe for deep/long graphs. |
+| `.backwards()` | **Primary.** Iterative topological sort — safe for deep/long graphs. Frees the graph as it runs (single-use; see the note under [Core concept](#core-concept-the-tensor)). |
 | `.backwards_recursive()` | Legacy recursive version; can hit Python's recursion limit on long sequences. Prefer `.backwards()`. |
 
 **Static helper**
@@ -231,7 +263,7 @@ Lookup table. Call with an integer index array; backward scatters gradients corr
 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`.
+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).
@@ -380,7 +412,7 @@ A full **encoder–decoder transformer** (the original *Attention Is All You Nee
 | `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.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. |
 
@@ -480,13 +512,41 @@ for step in range(num_steps):
     opt.clip_grads(1.0)
     opt.step()
 
-# inference: encode once, decode token by token, stop on EOS
+# 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.
+
+```python
+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`.
@@ -498,12 +558,21 @@ out = model.generate(src, sos_id=SOS, eos_id=EOS, max_new=32,
 - **`Sequential.parameters()` is shallow** — see the note under [Containers](#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 since the last release:
+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](#whats-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
@@ -514,14 +583,18 @@ Shipped since the last release:
 
 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)
-- Checkpoint save/load for the transformer models
 
 ---
 
 ## License
 
-MIT License. See [LICENSE](LICENSE) for details.
+APACHE License. See [LICENSE](LICENSE) for details.

{cogforge_engine-2.1.0 → cogforge_engine-2.1.2}/README.md

@@ -4,12 +4,15 @@
 
 `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.
+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](#whats-new).
 
 ---
 
 ## Table of contents
 
+- [What's new](#whats-new)
 - [Installation](#installation)
 - [Quick start](#quick-start)
 - [Backend: CPU, GPU, numexpr, and no-grad mode](#backend-cpu-gpu-numexpr-and-no-grad-mode)
@@ -27,12 +30,36 @@ There is no C++, no PyTorch — just NumPy and explicit, hand-derived gradients.
   - [Models](#models)
 - [Worked example: train a char-level GPT](#worked-example-train-a-char-level-gpt)
 - [Worked example: encoder–decoder Seq2Seq](#worked-example-encoderdecoder-seq2seq)
+- [Visualizing the computation graph](#visualizing-the-computation-graph)
 - [Gotchas](#gotchas)
 - [Roadmap](#roadmap)
 - [License](#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](#visualizing-the-computation-graph).
+
+---
+
 ## Installation
 
 ```bash
@@ -44,15 +71,17 @@ Requires Python 3.8+ and NumPy — that's the only hard dependency. Two optional
 ```bash
 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 three modules:
+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`). |
 
 ```python
 from cogforge.app import Tensor, Linear, Adam, MultiHeadAttention   # building blocks
@@ -143,9 +172,12 @@ Tensor(array, children=(), requires_grad=True, typed="compressed")
 | `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](#visualizing-the-computation-graph), draw it **before** calling `backwards()`.
+
 ---
 
 ## API reference
@@ -156,9 +188,9 @@ Gradients **accumulate** into `.grad`. Always zero them between optimization ste
 
 | Operation | Notes |
 | --- | --- |
-| `a + b`, `a - b`, `a * b` | Elementwise, with broadcasting support. |
+| `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. |
+| `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()` | |
@@ -176,7 +208,7 @@ Gradients **accumulate** into `.grad`. Always zero them between optimization ste
 
 | Method | Notes |
 | --- | --- |
-| `.backwards()` | **Primary.** Iterative topological sort — safe for deep/long graphs. |
+| `.backwards()` | **Primary.** Iterative topological sort — safe for deep/long graphs. Frees the graph as it runs (single-use; see the note under [Core concept](#core-concept-the-tensor)). |
 | `.backwards_recursive()` | Legacy recursive version; can hit Python's recursion limit on long sequences. Prefer `.backwards()`. |
 
 **Static helper**
@@ -216,7 +248,7 @@ Lookup table. Call with an integer index array; backward scatters gradients corr
 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`.
+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).
@@ -365,7 +397,7 @@ A full **encoder–decoder transformer** (the original *Attention Is All You Nee
 | `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.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. |
 
@@ -465,13 +497,41 @@ for step in range(num_steps):
     opt.clip_grads(1.0)
     opt.step()
 
-# inference: encode once, decode token by token, stop on EOS
+# 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.
+
+```python
+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`.
@@ -483,12 +543,21 @@ out = model.generate(src, sos_id=SOS, eos_id=EOS, max_new=32,
 - **`Sequential.parameters()` is shallow** — see the note under [Containers](#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 since the last release:
+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](#whats-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
@@ -499,14 +568,18 @@ Shipped since the last release:
 
 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)
-- Checkpoint save/load for the transformer models
 
 ---
 
 ## License
 
-MIT License. See [LICENSE](LICENSE) for details.
+APACHE License. See [LICENSE](LICENSE) for details.

{cogforge_engine-2.1.0 → cogforge_engine-2.1.2}/cogforge/app.py

@@ -857,49 +857,46 @@ class BatchNorm1D:
     
     def __call__(self, x: Tensor):
         reduce_dims = (0,) if x.data.ndim == 2 else (0, 1)
-        
+
         if self.training:
             mean = x.data.mean(axis=reduce_dims, keepdims=True)
-            var = x.data.var(axis=reduce_dims, keepdims=True)
-            x_centered = x.data - mean
+            var  = x.data.var(axis=reduce_dims, keepdims=True)
+
+            x_centered = x.data - mean          # local, not self.
             std_inv = 1.0 / backend.np.sqrt(var + self.eps)
             x_hat = x_centered * std_inv
-            
+
             N = numpy.prod([x.data.shape[d] for d in reduce_dims])
             unbiased_var = var.squeeze() * (N / (N - 1)) if N > 1 else var.squeeze()
-            
             self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.squeeze()
-            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbiased_var
-            
+            self.running_var  = (1 - self.momentum) * self.running_var  + self.momentum * unbiased_var
         else:
-            std_inv = 1.0 / backend.np.sqrt(self.running_var + self.eps)
-            x_hat = (x.data - self.running_mean) * std_inv
-        
+            std_inv = None
+            x_hat = (x.data - self.running_mean) / backend.np.sqrt(self.running_var + self.eps)
+
         out_data = self.gamma.data * x_hat + self.beta.data
         if backend.NO_GRAD: return Tensor(out_data, op="BatchNorm")
         out = Tensor(out_data, children=(x, self.gamma, self.beta), op="BatchNorm")
-        
+
+        was_training = self.training           
+
         def _backward():
-            if not self.training:
+            if not was_training:
                 return
-            
             dout = out.grad
             self.gamma.grad += backend.np.sum(dout * x_hat, axis=reduce_dims)
-            self.beta.grad += backend.np.sum(dout, axis=reduce_dims)
-            
+            self.beta.grad  += backend.np.sum(dout, axis=reduce_dims)
             N = numpy.prod([x.data.shape[d] for d in reduce_dims])
-
             dx_hat = dout * self.gamma.data
             dx = (1.0 / N) * std_inv * (
-                N * dx_hat 
-                - backend.np.sum(dx_hat, axis=reduce_dims, keepdims=True) 
+                N * dx_hat
+                - backend.np.sum(dx_hat, axis=reduce_dims, keepdims=True)
                 - x_hat * backend.np.sum(dx_hat * x_hat, axis=reduce_dims, keepdims=True)
+                # - x_hat * backend.np.sum(dx_hat * self.x_hat if False else dx_hat * x_hat, axis=reduce_dims, keepdims=True)
             )
-
             x.grad += dx
-            
+
         out._backwards = _backward
-        
         return out
             
            

{cogforge_engine-2.1.0 → cogforge_engine-2.1.2}/cogforge/models.py

@@ -233,8 +233,10 @@ class Seq2Seq:
         original_grad_state = not backend.NO_GRAD
         needGradientHence(False)
         try:
-            x_enc, enc_pad = self.encode(enc_idx)          # encode once
+            x_enc, enc_pad = self.encode(enc_idx)
             dec_idx = numpy.full((B, 1), sos_id, dtype=numpy.int64)
+            finished = numpy.zeros(B, dtype=bool)                    # NEW
+
             for _ in range(max_new):
                 logits = to_cpu(self.decode_step(dec_idx, x_enc, enc_pad).data)[:, -1, :] / temperature
                 if top_k is not None:
@@ -243,9 +245,14 @@ class Seq2Seq:
                 z = logits - logits.max(-1, keepdims=True)
                 p = numpy.exp(z).astype(numpy.float64); p /= p.sum(-1, keepdims=True)
                 nxt = numpy.array([[numpy.random.choice(len(pr), p=pr)] for pr in p])
+
+                nxt[finished, 0] = self.pad_id                       # NEW: freeze finished rows
                 dec_idx = numpy.concatenate([dec_idx, nxt], axis=1)
-                if eos_id is not None and bool((nxt[:, 0] == eos_id).all()):
-                    break
+
+                if eos_id is not None:
+                    finished |= (nxt[:, 0] == eos_id)                # NEW: mark newly finished
+                    if finished.all():                               # NEW: stop when everyone's done
+                        break
             return dec_idx
         finally:
             if was_training:

{cogforge_engine-2.1.0 → cogforge_engine-2.1.2}/cogforge/utils.py

@@ -1,4 +1,5 @@
 from graphviz import Digraph
+import numpy
 
 def trace_graph(root_tensor):
     """Recursively traces the tensor graph to find all nodes and edges."""
@@ -14,25 +15,43 @@ def trace_graph(root_tensor):
     build(root_tensor)
     return nodes, edges
 
-def draw_dot(root_tensor, format='svg', rankdir='LR'):
+def draw_graph(root_tensor, max_char=50,format='svg', rankdir='LR'):
     """
     Renders a computational graph starting from the given root_tensor.
     format: output format (e.g., 'png', 'svg', 'pdf')
     rankdir: 'LR' for left-to-right layout, 'TB' for top-to-bottom
     """
     dot = Digraph(format=format, graph_attr={'rankdir': rankdir})
-    
     nodes, edges = trace_graph(root_tensor)
     
+    def format_array(arr):
+        """Sanitizes and truncates array strings for safe Graphviz rendering."""
+        if arr is None:
+            return "None"
+        if isinstance(arr, numpy.ndarray):
+            s = numpy.array2string(arr, precision=4, suppress_small=True, separator=', ')
+        else:
+            s = str(arr)
+            
+        s = s.replace('\n', '')
+        
+        if len(s) > max_char:
+            s = s[:max_char-3] + "..."
+        s = s.replace('{', '\\{').replace('}', '\\}').replace('|', '\\|')
+        return s
+
     for n in nodes:
         uid = str(id(n))
         shape_str = str(n.shape)
         
+        data_str = format_array(n.data)
+        grad_str = format_array(n.grad)
+        
         if n.op:
-            label = f"{{ {n.op} | shape: {shape_str} }}"
+            label = f"{{ {n.op} | shape: {shape_str} | data: {data_str} | grad: {grad_str} }}"
             dot.node(name=uid, label=label, shape='record', style='filled', fillcolor='white')
         else:
-            label = f"{{ Data | shape: {shape_str} }}"
+            label = f"{{ Node | shape: {shape_str} | data: {data_str} | grad: {grad_str} }}"
             dot.node(name=uid, label=label, shape='record', style='filled', fillcolor='lightblue')
             
     for child, parent in edges:

{cogforge_engine-2.1.0 → cogforge_engine-2.1.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "cogforge-engine" 
-version = "2.1.0"
+version = "2.1.2"
 authors = [
   { name="Avik Majumder", email="avikmjd2@gmail.com" },
 ]