cogforge-engine 0.2.0__tar.gz → 0.2.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
File without changes
@@ -0,0 +1,330 @@
1
+ Metadata-Version: 2.4
2
+ Name: cogforge-engine
3
+ Version: 0.2.1
4
+ Summary: A custom autograd engine and Transformer block built from scratch.
5
+ Project-URL: Homepage, https://github.com/avikmjd2/cogforge
6
+ Author-email: Avik Majumder <avikmjd2@gmail.com>
7
+ License-File: LICENSE
8
+ Requires-Python: >=3.8
9
+ Requires-Dist: numpy>=1.20.0
10
+ Description-Content-Type: text/markdown
11
+
12
+ # cogforge
13
+
14
+ > A from-scratch deep learning library built on nothing but NumPy — a reverse-mode autograd engine extended all the way to a working GPT.
15
+
16
+ `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 attention, and a decoder-only transformer (`GPTV1`) you can actually train and sample from.
17
+
18
+ There is no C++, no CUDA, no PyTorch — just NumPy and explicit, hand-derived gradients. The goal is to *understand* every gradient that flows, not to be fast.
19
+
20
+ ---
21
+
22
+ ## Table of contents
23
+
24
+ - [Installation](#installation)
25
+ - [Quick start](#quick-start)
26
+ - [Core concept: the `Tensor`](#core-concept-the-tensor)
27
+ - [API reference](#api-reference)
28
+ - [Tensor — autograd engine](#tensor--autograd-engine)
29
+ - [Losses](#losses)
30
+ - [Layers](#layers)
31
+ - [Containers](#containers)
32
+ - [Recurrent](#recurrent)
33
+ - [Sequence-to-sequence](#sequence-to-sequence)
34
+ - [Optimizers](#optimizers)
35
+ - [Models](#models)
36
+ - [Worked example: train a char-level GPT](#worked-example-train-a-char-level-gpt)
37
+ - [Gotchas](#gotchas)
38
+ - [Roadmap](#roadmap)
39
+ - [License](#license)
40
+
41
+ ---
42
+
43
+ ## Installation
44
+
45
+ ```bash
46
+ pip install cogforge-engine
47
+ ```
48
+
49
+ Requires Python 3.8+ and NumPy. That's the only dependency.
50
+
51
+ The package is organized into two modules:
52
+
53
+ | Module | Contains |
54
+ | --- | --- |
55
+ | `cogforge.app` | The autograd engine (`Tensor`) and every building block — layers, optimizers, losses, normalization, attention. |
56
+ | `cogforge.models` | Ready-to-use models. Currently `GPTV1`. |
57
+
58
+ ```python
59
+ from cogforge.app import Tensor, Linear, Adam, MultiHeadAttention # building blocks
60
+ from cogforge.models import GPTV1 # models
61
+ ```
62
+
63
+ ---
64
+
65
+ ## Quick start
66
+
67
+ ```python
68
+ import numpy as np
69
+ from cogforge.app import Tensor
70
+
71
+ # Build a graph
72
+ a = Tensor(np.array([2.0, 3.0]))
73
+ b = Tensor(np.array([4.0, 5.0]))
74
+ c = (a * b).sigmoid().softmax()
75
+
76
+ # Backpropagate (note the spelling: backwards, with an 's')
77
+ c.backwards()
78
+
79
+ print(a.grad) # gradient of the output w.r.t. a
80
+ ```
81
+
82
+ Every `Tensor` carries a `.data` (the NumPy 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.
83
+
84
+ ---
85
+
86
+ ## Core concept: the `Tensor`
87
+
88
+ ```python
89
+ Tensor(array, children=(), requires_grad=True, typed="compressed")
90
+ ```
91
+
92
+ | Argument | Meaning |
93
+ | --- | --- |
94
+ | `array` | Any array-like; stored as a NumPy array in `.data`. |
95
+ | `children` | Parent tensors in the graph (set internally by ops; you rarely pass this). |
96
+ | `requires_grad` | Reserved flag (currently informational). |
97
+ | `typed` | `"compressed"` → `float32` (default), anything else → `float64`. |
98
+
99
+ Gradients **accumulate** into `.grad`. Always zero them between optimization steps (the optimizers do this for you via `zero_grad()`).
100
+
101
+ ---
102
+
103
+ ## API reference
104
+
105
+ ### Tensor — autograd engine
106
+
107
+ **Differentiable operations** (each builds graph and defines its own backward):
108
+
109
+ | Operation | Notes |
110
+ | --- | --- |
111
+ | `a + b`, `a - b`, `a * b` | Elementwise, with broadcasting support. |
112
+ | `a @ b` | Batched matmul; gradients are correctly un-broadcast. |
113
+ | `a[key]` | Indexing/slicing. |
114
+ | `.relu()` | |
115
+ | `.sigmoid()` | |
116
+ | `.tanh()` | |
117
+ | `.softmax(axis=-1)` | Numerically stable (max-subtraction). |
118
+ | `.view(shape)` | Reshape (handles non-contiguous data). |
119
+ | `.flatten()` | Flattens everything after the batch dim → `(B, -1)`. |
120
+ | `.flatten_consective(num)` | Groups `num` consecutive timesteps. Expects a 3-D `(B, T, C)` tensor; `T` must be divisible by `num`. |
121
+ | `.transpose(axes)` | Permute axes (pass the full permutation tuple). |
122
+ | `.masked_fill(mask, value)` | Sets entries where `mask` is `True` to `value` (used for causal attention). |
123
+
124
+ **Backward pass**
125
+
126
+ | Method | Notes |
127
+ | --- | --- |
128
+ | `.backwards()` | **Primary.** Iterative topological sort — safe for deep/long graphs. |
129
+ | `.backwards_recursive()` | Legacy recursive version; can hit Python's recursion limit on long sequences. Prefer `.backwards()`. |
130
+
131
+ **Static helper**
132
+
133
+ - `Tensor.unbroadcast(grad, shape)` — reduces a broadcasted gradient back to the original parameter shape. Used internally.
134
+
135
+ ---
136
+
137
+ ### Losses
138
+
139
+ 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.
140
+
141
+ | Loss | Input expectation | Use when |
142
+ | --- | --- | --- |
143
+ | `Tensor.cross_entropy_loss(predictions, targets)` | `predictions` are **probabilities** (call `.softmax()` first), `targets` one-hot. | You already have a softmax in your graph. |
144
+ | `Tensor.softmax_cross_entropy(scores, targets)` | `scores` are **raw logits**, `targets` one-hot. Softmax is fused inside (stable). Works for 2-D `(B,V)` and 3-D `(B,T,V)`. | Standard classification / LM. **Recommended.** |
145
+ | `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. |
146
+ | `Tensor.cross_entropy_loss_masked(predictions, targets, mask)` | Probabilities + per-row `mask` (1 = real, 0 = pad). | Padded batches. |
147
+ | `Tensor.softmax_cross_entropy_masked(scores, targets, mask)` | Logits + per-row mask. | Padded batches, fused softmax. |
148
+
149
+ > ℹ️ `softmax_cross_entropy` and `sparse_softmax_cross_entropy` apply softmax internally. Do **not** pass already-softmaxed values into them. `cross_entropy_loss` is the opposite — it expects probabilities.
150
+
151
+ ---
152
+
153
+ ### Layers
154
+
155
+ #### `Linear(nin, nout)`
156
+ Affine transform `x @ W + b`. He-initialized weights. `.parameters()` → `[W, b]`.
157
+
158
+ #### `Embedding(vocab_size, embedding_dim)`
159
+ Lookup table. Call with an integer index array; backward scatters gradients correctly (uses `np.add.at`, so repeated indices accumulate). `.parameters()` → `[weights]`.
160
+
161
+ #### `LayerNorm(dim, eps=1e-5)`
162
+ Normalizes over the last dimension. Learnable `gamma`/`beta`. `.parameters()` → `[gamma, beta]`.
163
+
164
+ #### `BatchNorm1D(dim, eps=1e-5, momentum=0.1)`
165
+ Normalizes over the batch (and time, for 3-D input). Tracks `running_mean`/`running_var` for inference. Toggle `.training = True/False`. Learnable `gamma`/`beta`.
166
+
167
+ #### `Attention(dk)`
168
+ Scaled dot-product attention. Call `attention(Q, K, V, mask=None)`. `dk` is the key dimension (sets the `1/√dk` scale).
169
+
170
+ #### `MultiHeadAttention(dinp, dmodel, dout, n)`
171
+ `n` heads, `dmodel` split into `n` chunks of size `dmodel // n` (must divide evenly). Projects input `dinp → dmodel`, attends, projects `dmodel → dout`. Call `mha(query, key, value, mask=None)`. `.parameters()` returns all four projection layers' params.
172
+
173
+ #### `FeedForward(dmodel, dff=None)`
174
+ Position-wise MLP: `Linear → ReLU → Linear`. `dff` defaults to `4 * dmodel`. `.parameters()` included.
175
+
176
+ #### `Transformer(dmodel, n, dff=None)`
177
+ A **pre-norm** decoder block: `x + Attn(LN(x))` then `x + FF(LN(x))`. `n` = number of attention heads. Call `block(x, mask=None)`. `.parameters()` included.
178
+
179
+ #### `PositionalEncoding(max_len, dmodel)`
180
+ Fixed sinusoidal positions, added to the input. Call `pe(x)`. No parameters.
181
+
182
+ ---
183
+
184
+ ### Containers
185
+
186
+ #### `Sequential(layers)`
187
+ Runs layers in order. `.train()` / `.test()` flip the `training` flag on any layer that has one (e.g. `BatchNorm1D`).
188
+
189
+ > ⚠️ `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()`.
190
+
191
+ #### `MLP(layer_sizes)`
192
+ Convenience feed-forward net: `Linear → ReLU` between layers, plain `Linear` output. Built from a list of sizes, e.g. `MLP([784, 128, 64, 10])`.
193
+
194
+ - `.save(filename="best_model.npz")` / `.load(filename="best_model.npz")` — persist/restore weights.
195
+ - *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.
196
+
197
+ ---
198
+
199
+ ### Recurrent
200
+
201
+ #### `RNNCell(input_dim, hidden_dim)`
202
+ One tanh recurrence step: `h_next = tanh(i2h(x) + h2h(h_prev))`. `.parameters()` included.
203
+
204
+ #### `RNN(input_dim, hidden_dim)`
205
+ 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`.
206
+
207
+ #### `StackedRNN(input_dim, hidden_dim, num_layers)`
208
+ Multiple `RNN` layers stacked. Returns `(top_layer_states, per_layer_final_states)` — the second value is convenient for seq2seq.
209
+
210
+ ---
211
+
212
+ ### Sequence-to-sequence
213
+
214
+ #### `Bridge(enc_hidden, dec_hidden, enc_layers, dec_layers, mode="project")`
215
+ Maps encoder final hidden states to decoder initial hidden states, handling mismatched layer counts and hidden sizes.
216
+
217
+ | `mode` | Behavior |
218
+ | --- | --- |
219
+ | `"project"` | One learned `Linear(enc_hidden → dec_hidden)` per decoder layer. General, recommended. |
220
+ | `"tie"` | No parameters; requires `enc_hidden == dec_hidden`. Selects/repeats raw states. |
221
+
222
+ ---
223
+
224
+ ### Optimizers
225
+
226
+ Both take an iterable of parameter tensors and share the same interface: `step()`, `zero_grad()`, `clip_grads(max_norm=5.0)`.
227
+
228
+ #### `SGD(parameters, learning_rate=0.01)`
229
+ Plain stochastic gradient descent.
230
+
231
+ #### `Adam(parameters, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8)`
232
+ Adam with bias correction. Recommended for transformers.
233
+
234
+ ```python
235
+ opt = Adam(model.parameters(), lr=3e-4)
236
+ opt.zero_grad()
237
+ loss.backwards()
238
+ opt.clip_grads(1.0) # optional gradient clipping
239
+ opt.step()
240
+ ```
241
+
242
+ ---
243
+
244
+ ### Models
245
+
246
+ #### `GPTV1(vocab, d_model, n_heads, n_layers, max_len, d_ff=None)`
247
+ A decoder-only transformer (token embedding + sinusoidal positions + stacked pre-norm `Transformer` blocks + final `LayerNorm` + output head). Causal masking is applied internally.
248
+
249
+ | Method | Description |
250
+ | --- | --- |
251
+ | `model(idx)` | `idx`: integer array `(B, T)`. Returns logits `(B, T, vocab)`. |
252
+ | `model.parameters()` | All trainable tensors. |
253
+ | `model.generate(idx, n_new, temperature=1.0, top_k=None)` | Autoregressive sampling. Crops to `max_len`, supports temperature and top-k. Returns `(B, T + n_new)`. |
254
+
255
+ ---
256
+
257
+ ## Worked example: train a char-level GPT
258
+
259
+ ```python
260
+ import numpy as np
261
+ from cogforge.app import Tensor, Adam
262
+ from cogforge.models import GPTV1
263
+
264
+ # --- data -------------------------------------------------------------
265
+ text = open("input.txt").read()
266
+ chars = sorted(set(text))
267
+ stoi = {c: i for i, c in enumerate(chars)}
268
+ itos = {i: c for i, c in enumerate(chars)}
269
+ data = np.array([stoi[c] for c in text])
270
+ vocab = len(chars)
271
+
272
+ # --- model ------------------------------------------------------------
273
+ block = 64
274
+ model = GPTV1(vocab=vocab, d_model=128, n_heads=4,
275
+ n_layers=4, max_len=block)
276
+ opt = Adam(model.parameters(), lr=3e-4)
277
+
278
+ def get_batch(bs=32):
279
+ ix = np.random.randint(0, len(data) - block - 1, size=bs)
280
+ x = np.stack([data[i:i + block] for i in ix])
281
+ y = np.stack([data[i + 1:i + block + 1] for i in ix])
282
+ return x, y
283
+
284
+ # --- train ------------------------------------------------------------
285
+ for step in range(2000):
286
+ x, y = get_batch()
287
+ logits = model(x) # (B, T, vocab)
288
+ loss = Tensor.sparse_softmax_cross_entropy(logits, y)
289
+
290
+ opt.zero_grad()
291
+ loss.backwards()
292
+ opt.clip_grads(1.0)
293
+ opt.step()
294
+
295
+ if step % 100 == 0:
296
+ print(f"step {step:4d} | loss {loss.data:.4f}")
297
+
298
+ # --- sample -----------------------------------------------------------
299
+ ctx = np.array([[stoi["\n"]]])
300
+ out = model.generate(ctx, n_new=300, temperature=0.8, top_k=20)
301
+ print("".join(itos[i] for i in out[0]))
302
+ ```
303
+
304
+ ---
305
+
306
+ ## Gotchas
307
+
308
+ - **It's `backwards()`, not `backward()`.** The backward pass method has a trailing `s`.
309
+ - **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.
310
+ - **Gradients accumulate.** Call `optimizer.zero_grad()` every step (or `p.grad[...] = 0`), or gradients pile up across iterations.
311
+ - **`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 `GPTV1.parameters()` does).
312
+ - **RNNs operate on lists**, not a single `(B, T, C)` tensor — pass a list of per-timestep tensors.
313
+
314
+ ---
315
+
316
+ ## Roadmap
317
+
318
+ Planned / under consideration:
319
+
320
+ - RoPE (rotary position embeddings) with length interpolation
321
+ - SwiGLU feed-forward and RMSNorm
322
+ - Weight tying between embedding and output head
323
+ - KV cache for faster generation
324
+ - Linear-attention block (as a study in the recall-vs-cost tradeoff)
325
+
326
+ ---
327
+
328
+ ## License
329
+
330
+ MIT License. See [LICENSE](LICENSE) for details.
@@ -0,0 +1,319 @@
1
+ # cogforge
2
+
3
+ > A from-scratch deep learning library built on nothing but NumPy — a reverse-mode autograd engine extended all the way to a working GPT.
4
+
5
+ `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 attention, and a decoder-only transformer (`GPTV1`) you can actually train and sample from.
6
+
7
+ There is no C++, no CUDA, no PyTorch — just NumPy and explicit, hand-derived gradients. The goal is to *understand* every gradient that flows, not to be fast.
8
+
9
+ ---
10
+
11
+ ## Table of contents
12
+
13
+ - [Installation](#installation)
14
+ - [Quick start](#quick-start)
15
+ - [Core concept: the `Tensor`](#core-concept-the-tensor)
16
+ - [API reference](#api-reference)
17
+ - [Tensor — autograd engine](#tensor--autograd-engine)
18
+ - [Losses](#losses)
19
+ - [Layers](#layers)
20
+ - [Containers](#containers)
21
+ - [Recurrent](#recurrent)
22
+ - [Sequence-to-sequence](#sequence-to-sequence)
23
+ - [Optimizers](#optimizers)
24
+ - [Models](#models)
25
+ - [Worked example: train a char-level GPT](#worked-example-train-a-char-level-gpt)
26
+ - [Gotchas](#gotchas)
27
+ - [Roadmap](#roadmap)
28
+ - [License](#license)
29
+
30
+ ---
31
+
32
+ ## Installation
33
+
34
+ ```bash
35
+ pip install cogforge-engine
36
+ ```
37
+
38
+ Requires Python 3.8+ and NumPy. That's the only dependency.
39
+
40
+ The package is organized into two modules:
41
+
42
+ | Module | Contains |
43
+ | --- | --- |
44
+ | `cogforge.app` | The autograd engine (`Tensor`) and every building block — layers, optimizers, losses, normalization, attention. |
45
+ | `cogforge.models` | Ready-to-use models. Currently `GPTV1`. |
46
+
47
+ ```python
48
+ from cogforge.app import Tensor, Linear, Adam, MultiHeadAttention # building blocks
49
+ from cogforge.models import GPTV1 # models
50
+ ```
51
+
52
+ ---
53
+
54
+ ## Quick start
55
+
56
+ ```python
57
+ import numpy as np
58
+ from cogforge.app import Tensor
59
+
60
+ # Build a graph
61
+ a = Tensor(np.array([2.0, 3.0]))
62
+ b = Tensor(np.array([4.0, 5.0]))
63
+ c = (a * b).sigmoid().softmax()
64
+
65
+ # Backpropagate (note the spelling: backwards, with an 's')
66
+ c.backwards()
67
+
68
+ print(a.grad) # gradient of the output w.r.t. a
69
+ ```
70
+
71
+ Every `Tensor` carries a `.data` (the NumPy 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.
72
+
73
+ ---
74
+
75
+ ## Core concept: the `Tensor`
76
+
77
+ ```python
78
+ Tensor(array, children=(), requires_grad=True, typed="compressed")
79
+ ```
80
+
81
+ | Argument | Meaning |
82
+ | --- | --- |
83
+ | `array` | Any array-like; stored as a NumPy array in `.data`. |
84
+ | `children` | Parent tensors in the graph (set internally by ops; you rarely pass this). |
85
+ | `requires_grad` | Reserved flag (currently informational). |
86
+ | `typed` | `"compressed"` → `float32` (default), anything else → `float64`. |
87
+
88
+ Gradients **accumulate** into `.grad`. Always zero them between optimization steps (the optimizers do this for you via `zero_grad()`).
89
+
90
+ ---
91
+
92
+ ## API reference
93
+
94
+ ### Tensor — autograd engine
95
+
96
+ **Differentiable operations** (each builds graph and defines its own backward):
97
+
98
+ | Operation | Notes |
99
+ | --- | --- |
100
+ | `a + b`, `a - b`, `a * b` | Elementwise, with broadcasting support. |
101
+ | `a @ b` | Batched matmul; gradients are correctly un-broadcast. |
102
+ | `a[key]` | Indexing/slicing. |
103
+ | `.relu()` | |
104
+ | `.sigmoid()` | |
105
+ | `.tanh()` | |
106
+ | `.softmax(axis=-1)` | Numerically stable (max-subtraction). |
107
+ | `.view(shape)` | Reshape (handles non-contiguous data). |
108
+ | `.flatten()` | Flattens everything after the batch dim → `(B, -1)`. |
109
+ | `.flatten_consective(num)` | Groups `num` consecutive timesteps. Expects a 3-D `(B, T, C)` tensor; `T` must be divisible by `num`. |
110
+ | `.transpose(axes)` | Permute axes (pass the full permutation tuple). |
111
+ | `.masked_fill(mask, value)` | Sets entries where `mask` is `True` to `value` (used for causal attention). |
112
+
113
+ **Backward pass**
114
+
115
+ | Method | Notes |
116
+ | --- | --- |
117
+ | `.backwards()` | **Primary.** Iterative topological sort — safe for deep/long graphs. |
118
+ | `.backwards_recursive()` | Legacy recursive version; can hit Python's recursion limit on long sequences. Prefer `.backwards()`. |
119
+
120
+ **Static helper**
121
+
122
+ - `Tensor.unbroadcast(grad, shape)` — reduces a broadcasted gradient back to the original parameter shape. Used internally.
123
+
124
+ ---
125
+
126
+ ### Losses
127
+
128
+ 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.
129
+
130
+ | Loss | Input expectation | Use when |
131
+ | --- | --- | --- |
132
+ | `Tensor.cross_entropy_loss(predictions, targets)` | `predictions` are **probabilities** (call `.softmax()` first), `targets` one-hot. | You already have a softmax in your graph. |
133
+ | `Tensor.softmax_cross_entropy(scores, targets)` | `scores` are **raw logits**, `targets` one-hot. Softmax is fused inside (stable). Works for 2-D `(B,V)` and 3-D `(B,T,V)`. | Standard classification / LM. **Recommended.** |
134
+ | `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. |
135
+ | `Tensor.cross_entropy_loss_masked(predictions, targets, mask)` | Probabilities + per-row `mask` (1 = real, 0 = pad). | Padded batches. |
136
+ | `Tensor.softmax_cross_entropy_masked(scores, targets, mask)` | Logits + per-row mask. | Padded batches, fused softmax. |
137
+
138
+ > ℹ️ `softmax_cross_entropy` and `sparse_softmax_cross_entropy` apply softmax internally. Do **not** pass already-softmaxed values into them. `cross_entropy_loss` is the opposite — it expects probabilities.
139
+
140
+ ---
141
+
142
+ ### Layers
143
+
144
+ #### `Linear(nin, nout)`
145
+ Affine transform `x @ W + b`. He-initialized weights. `.parameters()` → `[W, b]`.
146
+
147
+ #### `Embedding(vocab_size, embedding_dim)`
148
+ Lookup table. Call with an integer index array; backward scatters gradients correctly (uses `np.add.at`, so repeated indices accumulate). `.parameters()` → `[weights]`.
149
+
150
+ #### `LayerNorm(dim, eps=1e-5)`
151
+ Normalizes over the last dimension. Learnable `gamma`/`beta`. `.parameters()` → `[gamma, beta]`.
152
+
153
+ #### `BatchNorm1D(dim, eps=1e-5, momentum=0.1)`
154
+ Normalizes over the batch (and time, for 3-D input). Tracks `running_mean`/`running_var` for inference. Toggle `.training = True/False`. Learnable `gamma`/`beta`.
155
+
156
+ #### `Attention(dk)`
157
+ Scaled dot-product attention. Call `attention(Q, K, V, mask=None)`. `dk` is the key dimension (sets the `1/√dk` scale).
158
+
159
+ #### `MultiHeadAttention(dinp, dmodel, dout, n)`
160
+ `n` heads, `dmodel` split into `n` chunks of size `dmodel // n` (must divide evenly). Projects input `dinp → dmodel`, attends, projects `dmodel → dout`. Call `mha(query, key, value, mask=None)`. `.parameters()` returns all four projection layers' params.
161
+
162
+ #### `FeedForward(dmodel, dff=None)`
163
+ Position-wise MLP: `Linear → ReLU → Linear`. `dff` defaults to `4 * dmodel`. `.parameters()` included.
164
+
165
+ #### `Transformer(dmodel, n, dff=None)`
166
+ A **pre-norm** decoder block: `x + Attn(LN(x))` then `x + FF(LN(x))`. `n` = number of attention heads. Call `block(x, mask=None)`. `.parameters()` included.
167
+
168
+ #### `PositionalEncoding(max_len, dmodel)`
169
+ Fixed sinusoidal positions, added to the input. Call `pe(x)`. No parameters.
170
+
171
+ ---
172
+
173
+ ### Containers
174
+
175
+ #### `Sequential(layers)`
176
+ Runs layers in order. `.train()` / `.test()` flip the `training` flag on any layer that has one (e.g. `BatchNorm1D`).
177
+
178
+ > ⚠️ `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()`.
179
+
180
+ #### `MLP(layer_sizes)`
181
+ Convenience feed-forward net: `Linear → ReLU` between layers, plain `Linear` output. Built from a list of sizes, e.g. `MLP([784, 128, 64, 10])`.
182
+
183
+ - `.save(filename="best_model.npz")` / `.load(filename="best_model.npz")` — persist/restore weights.
184
+ - *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.
185
+
186
+ ---
187
+
188
+ ### Recurrent
189
+
190
+ #### `RNNCell(input_dim, hidden_dim)`
191
+ One tanh recurrence step: `h_next = tanh(i2h(x) + h2h(h_prev))`. `.parameters()` included.
192
+
193
+ #### `RNN(input_dim, hidden_dim)`
194
+ 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`.
195
+
196
+ #### `StackedRNN(input_dim, hidden_dim, num_layers)`
197
+ Multiple `RNN` layers stacked. Returns `(top_layer_states, per_layer_final_states)` — the second value is convenient for seq2seq.
198
+
199
+ ---
200
+
201
+ ### Sequence-to-sequence
202
+
203
+ #### `Bridge(enc_hidden, dec_hidden, enc_layers, dec_layers, mode="project")`
204
+ Maps encoder final hidden states to decoder initial hidden states, handling mismatched layer counts and hidden sizes.
205
+
206
+ | `mode` | Behavior |
207
+ | --- | --- |
208
+ | `"project"` | One learned `Linear(enc_hidden → dec_hidden)` per decoder layer. General, recommended. |
209
+ | `"tie"` | No parameters; requires `enc_hidden == dec_hidden`. Selects/repeats raw states. |
210
+
211
+ ---
212
+
213
+ ### Optimizers
214
+
215
+ Both take an iterable of parameter tensors and share the same interface: `step()`, `zero_grad()`, `clip_grads(max_norm=5.0)`.
216
+
217
+ #### `SGD(parameters, learning_rate=0.01)`
218
+ Plain stochastic gradient descent.
219
+
220
+ #### `Adam(parameters, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8)`
221
+ Adam with bias correction. Recommended for transformers.
222
+
223
+ ```python
224
+ opt = Adam(model.parameters(), lr=3e-4)
225
+ opt.zero_grad()
226
+ loss.backwards()
227
+ opt.clip_grads(1.0) # optional gradient clipping
228
+ opt.step()
229
+ ```
230
+
231
+ ---
232
+
233
+ ### Models
234
+
235
+ #### `GPTV1(vocab, d_model, n_heads, n_layers, max_len, d_ff=None)`
236
+ A decoder-only transformer (token embedding + sinusoidal positions + stacked pre-norm `Transformer` blocks + final `LayerNorm` + output head). Causal masking is applied internally.
237
+
238
+ | Method | Description |
239
+ | --- | --- |
240
+ | `model(idx)` | `idx`: integer array `(B, T)`. Returns logits `(B, T, vocab)`. |
241
+ | `model.parameters()` | All trainable tensors. |
242
+ | `model.generate(idx, n_new, temperature=1.0, top_k=None)` | Autoregressive sampling. Crops to `max_len`, supports temperature and top-k. Returns `(B, T + n_new)`. |
243
+
244
+ ---
245
+
246
+ ## Worked example: train a char-level GPT
247
+
248
+ ```python
249
+ import numpy as np
250
+ from cogforge.app import Tensor, Adam
251
+ from cogforge.models import GPTV1
252
+
253
+ # --- data -------------------------------------------------------------
254
+ text = open("input.txt").read()
255
+ chars = sorted(set(text))
256
+ stoi = {c: i for i, c in enumerate(chars)}
257
+ itos = {i: c for i, c in enumerate(chars)}
258
+ data = np.array([stoi[c] for c in text])
259
+ vocab = len(chars)
260
+
261
+ # --- model ------------------------------------------------------------
262
+ block = 64
263
+ model = GPTV1(vocab=vocab, d_model=128, n_heads=4,
264
+ n_layers=4, max_len=block)
265
+ opt = Adam(model.parameters(), lr=3e-4)
266
+
267
+ def get_batch(bs=32):
268
+ ix = np.random.randint(0, len(data) - block - 1, size=bs)
269
+ x = np.stack([data[i:i + block] for i in ix])
270
+ y = np.stack([data[i + 1:i + block + 1] for i in ix])
271
+ return x, y
272
+
273
+ # --- train ------------------------------------------------------------
274
+ for step in range(2000):
275
+ x, y = get_batch()
276
+ logits = model(x) # (B, T, vocab)
277
+ loss = Tensor.sparse_softmax_cross_entropy(logits, y)
278
+
279
+ opt.zero_grad()
280
+ loss.backwards()
281
+ opt.clip_grads(1.0)
282
+ opt.step()
283
+
284
+ if step % 100 == 0:
285
+ print(f"step {step:4d} | loss {loss.data:.4f}")
286
+
287
+ # --- sample -----------------------------------------------------------
288
+ ctx = np.array([[stoi["\n"]]])
289
+ out = model.generate(ctx, n_new=300, temperature=0.8, top_k=20)
290
+ print("".join(itos[i] for i in out[0]))
291
+ ```
292
+
293
+ ---
294
+
295
+ ## Gotchas
296
+
297
+ - **It's `backwards()`, not `backward()`.** The backward pass method has a trailing `s`.
298
+ - **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.
299
+ - **Gradients accumulate.** Call `optimizer.zero_grad()` every step (or `p.grad[...] = 0`), or gradients pile up across iterations.
300
+ - **`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 `GPTV1.parameters()` does).
301
+ - **RNNs operate on lists**, not a single `(B, T, C)` tensor — pass a list of per-timestep tensors.
302
+
303
+ ---
304
+
305
+ ## Roadmap
306
+
307
+ Planned / under consideration:
308
+
309
+ - RoPE (rotary position embeddings) with length interpolation
310
+ - SwiGLU feed-forward and RMSNorm
311
+ - Weight tying between embedding and output head
312
+ - KV cache for faster generation
313
+ - Linear-attention block (as a study in the recall-vs-cost tradeoff)
314
+
315
+ ---
316
+
317
+ ## License
318
+
319
+ MIT License. See [LICENSE](LICENSE) for details.
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "cogforge-engine"
7
- version = "0.2.0"
7
+ version = "0.2.1"
8
8
  authors = [
9
9
  { name="Avik Majumder", email="avikmjd2@gmail.com" },
10
10
  ]
@@ -1,11 +0,0 @@
1
- Metadata-Version: 2.4
2
- Name: cogforge-engine
3
- Version: 0.2.0
4
- Summary: A custom autograd engine and Transformer block built from scratch.
5
- Project-URL: Homepage, https://github.com/avikmjd2/cogforge
6
- Author-email: Avik Majumder <avikmjd2@gmail.com>
7
- Requires-Python: >=3.8
8
- Requires-Dist: numpy>=1.20.0
9
- Description-Content-Type: text/markdown
10
-
11
- Intial
@@ -1 +0,0 @@
1
- Intial