From f8b7e136d3632d267c3fc2fe001cf18881406ab6 Mon Sep 17 00:00:00 2001 From: "opencode-agent[bot]" Date: Sat, 1 Aug 2026 09:28:09 +0000 Subject: [PATCH 1/3] All 379 tests pass, 5-phase refactor done. Co-authored-by: sourcepirate --- docs/layers/base.md | 183 ++++++++---- docs/layers/core/merging.md | 163 ++++------ docs/models/model.md | 13 +- neutro/activations/base.py | 6 + neutro/activations/relu.py | 6 +- neutro/callbacks/base.py | 67 ++++- neutro/callbacks/checkpoint.py | 21 +- neutro/callbacks/early_stopping.py | 23 +- neutro/data.py | 21 +- neutro/engine/node.py | 16 +- neutro/initializers/base.py | 3 + neutro/layers/__init__.py | 6 + neutro/layers/base.py | 192 ++++++------ neutro/layers/core/dense.py | 10 +- neutro/layers/core/merging.py | 131 +++------ neutro/losses/base.py | 10 +- neutro/losses/mse.py | 4 +- neutro/metrics/base.py | 7 +- neutro/models/base_model.py | 457 +++++++++++------------------ neutro/models/moe/__init__.py | 0 neutro/models/vision/__init__.py | 3 + neutro/optimizers/adam.py | 35 ++- neutro/optimizers/adamw.py | 55 +--- neutro/optimizers/base.py | 14 +- neutro/optimizers/sgd.py | 18 +- neutro/tokenizers/bpe.py | 96 +++--- 26 files changed, 740 insertions(+), 820 deletions(-) delete mode 100644 neutro/models/moe/__init__.py diff --git a/docs/layers/base.md b/docs/layers/base.md index c4415ef..70a6792 100644 --- a/docs/layers/base.md +++ b/docs/layers/base.md @@ -70,50 +70,53 @@ def build(self, input_shape): ### Step 3: `__call__` — the dispatch hub -This is the most important method in the base class. It handles **two completely different modes** from a single entry point. +This is the most important method in the base class. It handles **two completely different modes** from a single entry point. The logic is factored into two helpers: `_symbolic_call` (graph-building) and `_eager_call` (computation). ```python def __call__(self, inputs, *args, **kwargs): - from ..engine.node import KerasTensor, Node + if self._is_symbolic_input(inputs): + return self._symbolic_call(inputs) + return self._eager_call(inputs, args, kwargs) +``` + +#### `_is_symbolic_input` + +```python +def _is_symbolic_input(self, inputs): + from ..engine.node import KerasTensor - is_symbolic = False if isinstance(inputs, KerasTensor): - is_symbolic = True - elif isinstance(inputs, list) and any(isinstance(i, KerasTensor) for i in inputs): - is_symbolic = True - - if is_symbolic: - # SYMBOLIC BRANCH — during model construction - if isinstance(inputs, list): - input_shapes = [i.shape for i in inputs] - else: - input_shapes = inputs.shape + return True + return (isinstance(inputs, list) + and any(isinstance(i, KerasTensor) for i in inputs)) +``` - if not self.built: - self.build(input_shapes) +🔍 **The fork**: If the input is a `KerasTensor` (or a list containing one), we're in "graph-building mode." If it's a real NumPy array, we're in "computation mode." - output_shape = self.compute_output_shape(input_shapes) +#### The symbolic branch — `_symbolic_call` - if isinstance(output_shape, list): - output_tensors = [KerasTensor(shape=s) for s in output_shape] - else: - output_tensors = KerasTensor(shape=output_shape) +```python +def _symbolic_call(self, inputs): + from ..engine.node import KerasTensor, Node - Node(self, input_tensors=inputs, output_tensors=output_tensors) - return output_tensors + if isinstance(inputs, list): + input_shapes = [i.shape for i in inputs] + else: + input_shapes = inputs.shape - # EAGER BRANCH — during training / inference if not self.built: - if isinstance(inputs, list): - self.build([i.shape for i in inputs]) - else: - self.build(inputs.shape) - return self.forward(inputs, *args, **kwargs) -``` + self.build(input_shapes) -🔍 **Lines 71-75**: `is_symbolic = ...` — The fork. If the input is a `KerasTensor` (or a list containing one), we're in "graph-building mode." If it's a real NumPy array, we're in "computation mode." + output_shape = self.compute_output_shape(input_shapes) -#### The symbolic branch (lines 77-97) + if isinstance(output_shape, list): + output_tensors = [KerasTensor(shape=s) for s in output_shape] + else: + output_tensors = KerasTensor(shape=output_shape) + + Node(self, input_tensors=inputs, output_tensors=output_tensors) + return output_tensors +``` When you use the Functional API like: @@ -124,17 +127,35 @@ x = Dense(64)(inputs) The `KerasTensor` called `inputs` is passed to `Dense.__call__`. No actual numbers flow through — just shape information. -🔍 **Lines 79-82**: `input_shapes = ...` — Extracts the shape from the symbolic tensor. Shapes look like `(None, 128)` where `None` means "unknown batch size." +🔍 **`input_shapes = ...`** — Extracts the shape from the symbolic tensor. Shapes look like `(None, 128)` where `None` means "unknown batch size." + +🔍 **`self.build(input_shapes)`** — Allocates weight matrices with the correct dimensions, but the actual *values* don't matter here. What matters is that `self.params['W']` now exists with the right shape. + +🔍 **`self.compute_output_shape(input_shapes)`** — Asks the layer: "If I give you input shape `(None, 128)`, what will my output shape be?" For a `Dense(64)` layer, the answer is `(None, 64)`. -🔍 **Line 84-85**: `self.build(input_shapes)` — Allocates weight matrices with the correct dimensions, but the actual *values* don't matter here. What matters is that `self.params['W']` now exists with the right shape. +🔍 **Creating output `KerasTensor`s** — Wraps the computed output shape into a new symbolic tensor. This tensor will be passed as input to the *next* layer. -🔍 **Line 87**: `self.compute_output_shape(input_shapes)` — Asks the layer: "If I give you input shape `(None, 128)`, what will my output shape be?" For a `Dense(64)` layer, the answer is `(None, 64)`. +🔍 **`Node(self, input_tensors=inputs, output_tensors=output_tensors)`** — Records the connection in the computation graph. This `Node` links "the input tensor(s)" to "the output tensor(s)" through "this layer." Later, `Model` walks these nodes to figure out the topology — which layers connect to which, what the forward pass order should be, and what the inputs/outputs of the whole model are. -🔍 **Lines 90-93**: Creating output `KerasTensor`s — Wraps the computed output shape into a new symbolic tensor. This tensor will be passed as input to the *next* layer. +#### The eager branch — `_eager_call` -🔍 **Line 96**: `Node(self, input_tensors=inputs, output_tensors=output_tensors)` — Records the connection in the computation graph. This `Node` links "the input tensor(s)" to "the output tensor(s)" through "this layer." Later, `Model` walks these nodes to figure out the topology — which layers connect to which, what the forward pass order should be, and what the inputs/outputs of the whole model are. +```python +def _eager_call(self, inputs, args, kwargs): + from neutro.autograd import as_tensor as _as_tensor -#### The eager branch (lines 99-105) + t_inputs = _as_tensor(inputs) + if not self.built: + if isinstance(t_inputs, list): + self.build([i.shape for i in t_inputs]) + else: + self.build(t_inputs.shape) + self._last_inputs = t_inputs + self._last_args = args + self._last_kwargs = kwargs + output = self.forward(t_inputs, *args, **kwargs) + self._last_output = output + return output +``` When you call a layer directly with real data: @@ -143,35 +164,40 @@ x = np.random.randn(32, 128) y = layer(x) # forwards! actual computation! ``` -🔍 **Lines 100-104**: `if not self.built: self.build(inputs.shape)` — First call? Build the weights using the actual concrete shape (e.g., `(32, 128)`). Note that `inputs.shape` here is a real tuple of integers, not a symbolic shape with `None`. +🔍 **`if not self.built: self.build(t_inputs.shape)`** — First call? Build the weights using the actual concrete shape (e.g., `(32, 128)`). Note that `inputs.shape` here is a real tuple of integers, not a symbolic shape with `None`. + +🔍 **`self._last_inputs = t_inputs`** — Caches the input for the `backward` method, which re-runs the forward pass inside a `GradientTape` to compute gradients. -🔍 **Line 105**: `return self.forward(inputs, *args, **kwargs)` — Delegates to the subclass's actual computation. This is where the matrix multiply happens, where the convolution runs, where the attention scores are computed. +🔍 **`return self.forward(t_inputs, *args, **kwargs)`** — Delegates to the subclass's actual computation. This is where the matrix multiply happens, where the convolution runs, where the attention scores are computed. ### Step 4: `sublayers` — finding nested layers ```python @property def sublayers(self): + """Return all nested layers contained within this layer. + + Traverses instance attributes, including sublayers stored in lists, + tuples, or dictionaries, with cycle protection and de-duplication. + """ layers = [] - for attr_name in dir(self): - if attr_name.startswith('_') or attr_name == 'sublayers': - continue - try: - attr = getattr(self, attr_name) - except AttributeError: - continue + visited_ids = {id(self)} + stack = [] + for attr in vars(self).values(): + stack.append(attr) + while stack: + attr = stack.pop() if isinstance(attr, Layer): - layers.append(attr) - elif isinstance(attr, list): - stack = [attr] - while stack: - curr = stack.pop() - for item in curr: - if isinstance(item, Layer): - layers.append(item) - elif isinstance(item, list): - stack.append(item) + if id(attr) not in visited_ids: + visited_ids.add(id(attr)) + layers.append(attr) + for nested in vars(attr).values(): + stack.append(nested) + elif isinstance(attr, (list, tuple, set)): + stack.extend(attr) + elif isinstance(attr, dict): + stack.extend(attr.values()) return layers ``` @@ -186,12 +212,12 @@ class TransformerBlock(Layer): When the optimizer needs to find **all** trainable parameters, it calls `sublayers` on the top-level model. The property: -1. Iterates over every attribute of the layer using `dir(self)` — this includes attributes defined in `__init__` of the current class **and** parent classes. -2. Skips private attributes (starting with `_`) and the property itself (to avoid infinite recursion). -3. If an attribute is a `Layer` instance, it collects it — this catches `self.attention`, `self.norm`, etc. -4. If an attribute is a **list**, it recursively searches inside it — this catches `self.ffn = [Dense(512), Dense(512)]`. It even handles lists-of-lists (used by `MoELayer` which has a list of expert lists). +1. Collects every instance attribute of the layer (via `vars(self)`) into a work stack. +2. If an attribute is a `Layer` instance, it collects it and pushes that layer's own attributes onto the stack — this catches `self.attention`, `self.norm`, etc. The `visited_ids` set de-duplicates, so a shared layer used in multiple branches is only reported once. +3. If an attribute is a **list, tuple, or set**, it pushes each element — this catches `self.ffn = [Dense(512), Dense(512)]` and the lists-of-lists used by `MoELayer` (a list of expert lists). +4. If an attribute is a **dict**, it pushes the values. -🔍 **Why is this important?** Without `sublayers`, a `TransformerBlock` would report only its own `params` dict (which is empty — it delegates everything to sublayers). With `sublayers`, the optimizer can traverse the full hierarchy and find every weight matrix in every attention head and every feed-forward layer. +🔍 **Why is this important?** Without `sublayers`, a `TransformerBlock` would report only its own `params` dict (which is empty — it delegates everything to sublayers). With `sublayers`, the optimizer can traverse the full hierarchy and find every weight matrix in every attention head and every feed-forward layer. The cycle protection (`visited_ids`) and de-duplication prevent infinite recursion and double-counting when layers are shared. ### Step 5: `count_params` — the recursive parameter counter @@ -218,14 +244,43 @@ def compute_output_shape(self, input_shape): return input_shape ``` -🔍 **Lines 55-62**: Default behavior — if no `output_shape` was explicitly set, assume the output shape equals the input shape. Subclasses like `Dense` override this to return `(*input_shape[:-1], units)`. +🔍 **Default behavior** — if no `output_shape` was explicitly set, assume the output shape equals the input shape. Subclasses like `Dense` override this to return `(*input_shape[:-1], units)`. + +#### `backward` + +The base `Layer.backward` provides a **generic gradient mechanism** for layers whose `forward` is written with autograd `Tensor` operations. It re-runs the forward pass inside a `GradientTape` to recover the parameter gradients via reverse-mode differentiation: ```python def backward(self, grad_output): - raise NotImplementedError + from neutro.autograd import GradientTape, Tensor as AT + + t_inputs = self._captured_inputs() # autograd view of the last inputs + if t_inputs is None: + return np.asarray(grad_output) + + fwd_args = getattr(self, '_last_args', ()) + fwd_kwargs = getattr(self, '_last_kwargs', {}).copy() + fwd_kwargs.pop('kv_cache', None) # runtime bookkeeping, not differentiable + fwd_kwargs.pop('layer_id', None) + + all_sources = list(self._collect_tensor_params().values()) + # ... include t_inputs in all_sources ... + + with GradientTape() as tape: + for source in all_sources: + tape.watch(source) + output = self.forward(t_inputs, *fwd_args, **fwd_kwargs) + g_t = AT(np.asarray(grad_output)) + loss = (output * g_t).sum() # dot product = weighted gradient probe + + tape.gradient(loss, all_sources) + # ... copy param_value.grad into layer.grads for every sublayer ... + return t_inputs.grad.copy() # gradient w.r.t. the inputs ``` -🔍 **Line 64-65**: The base class doesn't know how to backpropagate (that depends on the concrete computation). Subclasses **must** implement this. If they don't, calling `backward` will crash with `NotImplementedError` — a clear signal that you forgot to implement it. +🔍 **Why re-run forward?** The gradient of the weighted probe `loss = (output * g_t).sum()` w.r.t. any parameter is exactly the upstream gradient `g_t` propagated through the layer's computation graph. This gives a correct, dependency-free backward pass for any forward written purely with autograd ops. + +🔍 **`_captured_inputs` / `_collect_tensor_params`** — helper methods that rebuild the input `Tensor`s and gather every autograd parameter across the layer tree (including sublayers). This is what makes `backward` work for nested layers like `TransformerBlock` without each sublayer reimplementing graph traversal. ## Putting it all together diff --git a/docs/layers/core/merging.md b/docs/layers/core/merging.md index 57ad22e..90c959d 100644 --- a/docs/layers/core/merging.md +++ b/docs/layers/core/merging.md @@ -2,7 +2,23 @@ Merge layers combine **multiple input tensors** into a single output tensor. They are essential for building non-linear architectures like ResNets (skip connections), Inception modules, and multi-branch networks. Every merge layer takes a **list of tensors** as input. -## Add — `merging.py:4` +## Shared structure + +All reduce-type merge layers (`Add`, `Multiply`, `Average`, `Maximum`, `Minimum`) inherit from `_ReduceBase` (`merging.py:11`), which provides the common `build`, `compute_output_shape`, and a `forward` that folds a single binary operation over the input list. Each concrete layer only implements `_combine(left, right)` — the per-pair reduction: + +```python +def forward(self, inputs, training=False): + if not isinstance(inputs, list): + return inputs + result = _to_tensor(inputs[0]) + for other in inputs[1:]: + result = self._combine(result, _to_tensor(other)) + return result +``` + +`_to_tensor` coerces plain NumPy inputs into autograd `Tensor`s so the whole reduction is differentiable. + +## Add — `merging.py:42` ### What does this layer do? @@ -18,32 +34,20 @@ Every input must have the **same shape**. The output has that same shape. ### Walking through the code -#### `forward` +#### `_combine` ```python -def forward(self, inputs, training=False): - self.input_lengths = len(inputs) - return sum(inputs) +def _combine(self, left, right): + return left + right ``` -🔍 **Line `self.input_lengths = len(inputs)`**: We cache the number of inputs. The backward pass needs this to know how many gradient tensors to return. - -🔍 **Line `return sum(inputs)`**: Python's built-in `sum()` on a list of NumPy arrays performs element-wise addition. All arrays must have the same shape. For example, listing `[a, b, c]` computes `a + b + c`. - 📐 **Shape**: If each input is `(batch, 64)`, the output is also `(batch, 64)`. -#### `backward` - -```python -def backward(self, grad_output): - return [grad_output for _ in range(self.input_lengths)] -``` - -🔍 **Line `[grad_output for _ in range(self.input_lengths)]`**: For $y = x_1 + x_2$, we have $\partial y / \partial x_1 = 1$ and $\partial y / \partial x_2 = 1$. So by the chain rule, $\partial L / \partial x_i = \partial L / \partial y \cdot 1$. The gradient is **broadcast unchanged** to every input. We return a list with `N` identical gradient tensors. +The additive reduction routes the gradient unchanged back to every input, so skip connections receive clean gradient flow. --- -## Concatenate — `merging.py:42` +## Concatenate — `merging.py:47` ### What does this layer do? @@ -123,7 +127,7 @@ def backward(self, grad_output): --- -## Multiply — `merging.py:85` +## Multiply — `merging.py:79` ### What does this layer do? @@ -145,55 +149,35 @@ $$ ### Walking through the code -#### `forward` +#### `_combine` ```python -def forward(self, inputs, training=False): - self.inputs = inputs - res = inputs[0].copy() - for i in range(1, len(inputs)): - res *= inputs[i] - return res +def _combine(self, left, right): + return left * right ``` -🔍 **Line `self.inputs = inputs`**: Cache the list of inputs for the backward pass. The backward pass needs to access all inputs except the one being differentiated. - -🔍 **Line `res = inputs[0].copy()`**: Start with a **copy** of the first input. We use `.copy()` to avoid mutating the original input tensor. - -🔍 **Line `res *= inputs[i]`**: Multiply element-by-element. After the loop, `res` is the product of all inputs. +The base `_ReduceBase.forward` folds this pairwise product over all inputs: `(a * b) * c`. 📐 **Shape**: `(8, 64)` × `(8, 64)` × `(8, 64)` → `(8, 64)`. #### `backward` -```python -def backward(self, grad_output): - grads = [] - for i in range(len(self.inputs)): - g = grad_output.copy() - for j in range(len(self.inputs)): - if i == j: - continue - g *= self.inputs[j] - grads.append(g) - return grads -``` - -🔍 **Line `g = grad_output.copy()`**: Start with the upstream gradient. +The `Multiply` reduction is performed on autograd `Tensor`s, so the reverse-mode engine (`GradientTape`) computes the gradients automatically. For $y = x_1 \odot x_2 \odot \cdots \odot x_N$: -🔍 **Line `for j ... if i == j: continue; g *= self.inputs[j]`**: -For input $x_i$, we multiply the upstream gradient by **every other input** $x_j$ for $j \neq i$. This implements $\partial L / \partial x_i = \partial L / \partial y \cdot \prod_{j \neq i} x_j$. +$$ +\frac{\partial L}{\partial x_i} = \frac{\partial L}{\partial y} \odot \prod_{j \neq i} x_j +$$ 📐 **Example with 3 inputs**: $y = a \cdot b \cdot c$. - $\partial L / \partial a = \partial L / \partial y \cdot b \cdot c$ - $\partial L / \partial b = \partial L / \partial y \cdot a \cdot c$ - $\partial L / \partial c = \partial L / \partial y \cdot a \cdot b$ -The loops compute exactly these products. +The autograd engine computes exactly these products. --- -## Average — `merging.py:120` +## Average — `merging.py:84` ### What does this layer do? @@ -211,26 +195,21 @@ $$ ```python def forward(self, inputs, training=False): - self.input_lengths = len(inputs) - return sum(inputs) / self.input_lengths + result = super().forward(inputs, training=training) # sum via _combine + if isinstance(inputs, list) and len(inputs) > 1: + result = result / float(len(inputs)) + return result ``` -🔍 **Line `self.input_lengths = len(inputs)`**: Cache the number of inputs `N` for the backward pass. - -🔍 **Line `sum(inputs) / self.input_lengths`**: Python's `sum()` adds element-wise, then dividing by `N` gives the average. +The base `_ReduceBase.forward` sums all inputs with `_combine(a, b) = a + b`, then `Average` divides by the number of inputs `N`. #### `backward` -```python -def backward(self, grad_output): - return [grad_output / self.input_lengths for _ in range(self.input_lengths)] -``` - -🔍 **Line `grad_output / self.input_lengths`**: The derivative of $y = (x_1 + \dots + x_N) / N$ w.r.t. $x_i$ is $1/N$. Each input receives the upstream gradient divided by the number of inputs. +The derivative of $y = (x_1 + \dots + x_N) / N$ w.r.t. $x_i$ is $1/N$. Each input receives the upstream gradient divided by the number of inputs, computed automatically by the autograd engine. --- -## Maximum — `merging.py:144` +## Maximum — `merging.py:95` ### What does this layer do? @@ -248,51 +227,31 @@ The backward pass uses **argmax routing**: the gradient flows only to the input( ### Walking through the code -#### `forward` +#### `_combine` ```python -def forward(self, inputs, training=False): - self.inputs = inputs - res = inputs[0].copy() - for i in range(1, len(inputs)): - res = np.maximum(res, inputs[i]) - return res +def _combine(self, left, right): + return autograd_ops.maximum(left, right) ``` -🔍 **Line `self.inputs = inputs`**: Cache the inputs. The backward pass needs to compare each input against the maximum. - -🔍 **Line `np.maximum(res, inputs[i])`**: Element-wise maximum. `np.maximum(a, b)` returns an array where each element is `max(a_element, b_element)`. +The base `_ReduceBase.forward` folds this pairwise `maximum` over all inputs. `autograd_ops.maximum` is differentiable, so the backward pass runs through the autograd engine. 📐 **Shape**: All `(8, 64)`. Output: `(8, 64)`. #### `backward` -```python -def backward(self, grad_output): - max_val = self.forward(self.inputs) - grads = [] - for inp in self.inputs: - mask = (inp == max_val) - grads.append(grad_output * mask) - return grads -``` - -🔍 **Line `max_val = self.forward(self.inputs)`**: Recompute the maximum values by calling `forward` again. (Alternative: cache `max_val` in forward.) +The autograd engine uses **argmax routing**: the gradient flows only to the input(s) that actually **were** the maximum at each position. All other inputs receive zero gradient. -🔍 **Line `mask = (inp == max_val)`**: For each input, create a boolean mask that is `True` wherever this input equals the maximum value. If multiple inputs share the maximum at a position, all of them get gradient. - -🔍 **Line `grad_output * mask`**: The mask zeros out the gradient everywhere this input was **not** the maximum. Only the "winning" input receives gradient. - -📐 **The logic**: For $y = \max(x_1, x_2)$, the subgradient is: +🔍 **The logic**: For $y = \max(x_1, x_2)$, the subgradient is: $$ \frac{\partial y}{\partial x_1} = \begin{cases} 1 & \text{if } x_1 > x_2 \\ 0 & \text{if } x_1 < x_2 \\ \text{any value in } [0,1] & \text{if } x_1 = x_2 \end{cases} $$ -Neutro uses the tie-case convention: if two inputs are equal, **both** get gradient (the mask is `True` for both). +Neutro uses the tie-case convention: if two inputs are equal, **both** get gradient. --- -## Minimum — `merging.py:177` +## Minimum — `merging.py:100` ### What does this layer do? @@ -308,34 +267,18 @@ The backward pass uses **argmin routing**: gradient flows only to the input(s) t ### Walking through the code -#### `forward` +#### `_combine` ```python -def forward(self, inputs, training=False): - self.inputs = inputs - res = inputs[0].copy() - for i in range(1, len(inputs)): - res = np.minimum(res, inputs[i]) - return res +def _combine(self, left, right): + return -autograd_ops.maximum(-left, -right) ``` -Identical to Maximum but uses `np.minimum`. +`Minimum` is implemented as the negation of `Maximum`: `min(a, b) = -max(-a, -b)`. This reuses the differentiable `maximum` op. #### `backward` -```python -def backward(self, grad_output): - min_val = self.forward(self.inputs) - grads = [] - for inp in self.inputs: - mask = (inp == min_val) - grads.append(grad_output * mask) - return grads -``` - -Identical to Maximum's backward but using the minimum value as the comparison target. - -🔍 **Line `mask = (inp == min_val)`**: Gradient passes only where this input equals the minimum. For ties, multiple inputs receive gradient. +Identical to Maximum's backward: gradient passes only where this input equals the minimum (argmin routing), computed by the autograd engine. For ties, multiple inputs receive gradient. --- diff --git a/docs/models/model.md b/docs/models/model.md index 51fe891..69113a9 100644 --- a/docs/models/model.md +++ b/docs/models/model.md @@ -149,7 +149,7 @@ def _capture_layer_state(layer): l = stack.pop() if id(l) in visited: continue visited.add(id(l)) - sub = {k: v for k, v in l.__dict__.items() + sub = {k: copy.deepcopy(v) for k, v in l.__dict__.items() if k not in Model._STATE_EXCLUDE} state[id(l)] = sub for sl in l.sublayers: @@ -157,17 +157,24 @@ def _capture_layer_state(layer): return state ``` -This recursively captures the `__dict__` of every sublayer, keyed by `id()`. Excluded keys (`params`, `grads`, `built`, `input_shape`, etc.) are persistent architectural attributes that should not be restored. +This recursively captures the `__dict__` of every sublayer, keyed by `id()`. Values are **deep-copied** so that mutations during later forward passes (e.g. a shared layer used in multiple branches) cannot corrupt a previously captured snapshot. Excluded keys (`params`, `grads`, `built`, `input_shape`, etc.) are persistent architectural attributes that should not be restored. ### The `fit` Method -Supports three input modes: +`fit` and `autograd_fit` share a single training loop (`Model._fit_loop`). The only difference is the per-batch step: + +- **`fit`** (default) — uses manual backward passes (`layer.backward`) through the `Model.backward` graph traversal. +- **`autograd_fit`** — uses the autograd engine: a `GradientTape` watches `trainable_params`, and gradients are produced by `tape.gradient(loss, params)`. + +Both are thin wrappers over `_fit_loop`, which supports three input modes: 1. **Single array**: `fit(x, y)` — standard training. 2. **List of arrays (MIMO)**: `fit([x1, x2], [y1, y2])` — multi-input, multi-output. - `is_mimo_x = isinstance(x, list)` detects list inputs. - For MIMO, batch slicing uses `[xi[start:end] for xi in x_shuffled]`. 3. **Generator**: `fit(generator)` — yields `(x_batch, y_batch)` tuples. +The per-batch logic lives in `Model._train_on_batch(x_batch, y_batch, use_autograd)`, which runs one forward/backward/optimizer step and returns the scalar batch loss plus the model output (used for metric reporting). + Loss is summed across multiple outputs (matching Keras behavior): `batch_loss = sum(self.loss_fn(y_batch[j], output[j])`. ### `evaluate` diff --git a/neutro/activations/base.py b/neutro/activations/base.py index 2d5150d..acbaf4d 100644 --- a/neutro/activations/base.py +++ b/neutro/activations/base.py @@ -1,7 +1,13 @@ from neutro.autograd import Tensor as AutogradTensor + class Activation: + """Base class for activation functions.""" + def __call__(self, x): + """Apply the activation function.""" raise NotImplementedError + def gradient(self, x): + """Compute the derivative of the activation function.""" raise NotImplementedError diff --git a/neutro/activations/relu.py b/neutro/activations/relu.py index 16c592b..a6f1059 100644 --- a/neutro/activations/relu.py +++ b/neutro/activations/relu.py @@ -2,8 +2,12 @@ from .base import Activation from neutro.autograd import ops as autograd_ops + class ReLU(Activation): + """Rectified Linear Unit activation.""" + def __call__(self, x): return autograd_ops.relu(x) + def gradient(self, x): - return (x > 0).astype(float) \ No newline at end of file + return (x > 0).astype(float) diff --git a/neutro/callbacks/base.py b/neutro/callbacks/base.py index fa8fdb2..67dd60c 100644 --- a/neutro/callbacks/base.py +++ b/neutro/callbacks/base.py @@ -1,13 +1,62 @@ +from typing import Any, Dict, Optional + + class Callback: - def __init__(self): - self.model = None + """Base class for training callbacks. + + Subclasses override the lifecycle hooks below to observe or modify + training. `set_model` is called before training begins. + """ + + def __init__(self) -> None: + self.model: Optional[Any] = None - def set_model(self, model): + def set_model(self, model) -> None: self.model = model - def on_epoch_begin(self, epoch, logs=None): pass - def on_epoch_end(self, epoch, logs=None): pass - def on_batch_begin(self, batch, logs=None): pass - def on_batch_end(self, batch, logs=None): pass - def on_train_begin(self, logs=None): pass - def on_train_end(self, logs=None): pass + def on_train_begin(self, logs=None) -> None: + pass + + def on_train_end(self, logs=None) -> None: + pass + + def on_epoch_begin(self, epoch, logs=None) -> None: + pass + + def on_epoch_end(self, epoch, logs=None) -> None: + pass + + def on_batch_begin(self, batch, logs=None) -> None: + pass + + def on_batch_end(self, batch, logs=None) -> None: + pass + + +class MonitorCallback(Callback): + """Base class for callbacks that track an epoch-level metric. + + Provides shared logic for deciding whether a monitored metric has + improved, mirroring the Keras ``monitor``/``mode`` semantics. + """ + + def __init__(self, monitor: str = 'val_loss', mode: str = 'auto') -> None: + super().__init__() + self.monitor = monitor + self.mode = mode + + def _is_improvement(self, current: float) -> bool: + if self.mode == 'min': + return current < self.best + if self.mode == 'max': + return current > self.best + # auto: infer direction from the monitored metric's name + if 'acc' in self.monitor: + return current > self.best + if 'loss' in self.monitor: + return current < self.best + return current < self.best + + def _init_best(self) -> float: + maximize = self.mode == 'max' or (self.mode == 'auto' and 'acc' in self.monitor) + return -float('inf') if maximize else float('inf') diff --git a/neutro/callbacks/checkpoint.py b/neutro/callbacks/checkpoint.py index 6733eb5..06e0c9f 100644 --- a/neutro/callbacks/checkpoint.py +++ b/neutro/callbacks/checkpoint.py @@ -1,24 +1,23 @@ -import numpy as np -from .base import Callback +from .base import MonitorCallback + + +class ModelCheckpoint(MonitorCallback): + """Save the model to disk after each epoch (optionally only on improvement).""" -class ModelCheckpoint(Callback): def __init__(self, filepath, monitor='val_loss', save_best_only=False, mode='auto'): - super().__init__() + super().__init__(monitor=monitor, mode=mode) self.filepath = filepath - self.monitor = monitor self.save_best_only = save_best_only - self.best = -np.inf if mode == 'max' or (mode == 'auto' and 'acc' in monitor) else np.inf - self.mode = mode + self.best = self._init_best() def on_epoch_end(self, epoch, logs=None): logs = logs or {} current = logs.get(self.monitor) - if current is None: return + if current is None: + return if self.save_best_only: - if (self.mode == 'min' and current < self.best) or \ - (self.mode == 'max' and current > self.best) or \ - (self.mode == 'auto' and (('acc' in self.monitor and current > self.best) or ('loss' in self.monitor and current < self.best))): + if self._is_improvement(current): self.best = current self.model.save(self.filepath) else: diff --git a/neutro/callbacks/early_stopping.py b/neutro/callbacks/early_stopping.py index c26f961..dbd24ab 100644 --- a/neutro/callbacks/early_stopping.py +++ b/neutro/callbacks/early_stopping.py @@ -1,27 +1,26 @@ -import numpy as np -from .base import Callback +from .base import MonitorCallback + + +class EarlyStopping(MonitorCallback): + """Stop training when a monitored metric has stopped improving.""" -class EarlyStopping(Callback): def __init__(self, monitor='val_loss', patience=0, mode='auto'): - super().__init__() - self.monitor = monitor + super().__init__(monitor=monitor, mode=mode) self.patience = patience self.wait = 0 - self.best = -np.inf if mode == 'max' or (mode == 'auto' and 'acc' in monitor) else np.inf - self.mode = mode + self.best = self._init_best() def on_epoch_end(self, epoch, logs=None): logs = logs or {} current = logs.get(self.monitor) - if current is None: return + if current is None: + return - if (self.mode == 'min' and current < self.best) or \ - (self.mode == 'max' and current > self.best) or \ - (self.mode == 'auto' and (('acc' in self.monitor and current > self.best) or ('loss' in self.monitor and current < self.best))): + if self._is_improvement(current): self.best = current self.wait = 0 else: self.wait += 1 if self.wait >= self.patience: self.model.stop_training = True - print(f"Epoch {epoch+1}: early stopping") + print(f"Epoch {epoch + 1}: early stopping") diff --git a/neutro/data.py b/neutro/data.py index 801aeb4..43b6b2d 100644 --- a/neutro/data.py +++ b/neutro/data.py @@ -1,16 +1,19 @@ import numpy as np + class DataLoader: - """ - Data loader for batching and shuffling data. - + """Data loader for batching and shuffling data. + Args: x: Input data (NumPy array). y: Target data (NumPy array). batch_size: Number of samples per batch. shuffle: Whether to shuffle the data at the beginning of each epoch. + augmenter: Optional object with an ``apply_transform`` method. """ - def __init__(self, x, y, batch_size=32, shuffle=True, augmenter=None): + + def __init__(self, x, y, batch_size: int = 32, shuffle: bool = True, + augmenter=None) -> None: self.x = x self.y = y self.batch_size = batch_size @@ -19,23 +22,23 @@ def __init__(self, x, y, batch_size=32, shuffle=True, augmenter=None): self.indices = np.arange(len(x)) self.on_epoch_end() - def __len__(self): + def __len__(self) -> int: return int(np.ceil(len(self.x) / self.batch_size)) - def on_epoch_end(self): + def on_epoch_end(self) -> None: if self.shuffle: np.random.shuffle(self.indices) def __getitem__(self, index): - indices = self.indices[index * self.batch_size : (index + 1) * self.batch_size] + indices = self.indices[index * self.batch_size:(index + 1) * self.batch_size] batch_x, batch_y = self.x[indices], self.y[indices] - + if self.augmenter: augmented_x = np.zeros_like(batch_x) for i in range(len(batch_x)): augmented_x[i] = self.augmenter.apply_transform(batch_x[i]) batch_x = augmented_x - + return batch_x, batch_y def __iter__(self): diff --git a/neutro/engine/node.py b/neutro/engine/node.py index 7cad2f2..3ba2c4c 100644 --- a/neutro/engine/node.py +++ b/neutro/engine/node.py @@ -1,9 +1,9 @@ import numpy as np + class KerasTensor: - """ - Symbolic representation of a tensor in the functional API. - """ + """Symbolic representation of a tensor in the functional API.""" + def __init__(self, shape, node=None, name=None): self.shape = shape self.node = node # The node that produced this tensor @@ -12,21 +12,23 @@ def __init__(self, shape, node=None, name=None): def __repr__(self): return f"KerasTensor(shape={self.shape}, name={self.name})" + class Node: - """ - Represents a 'call' to a layer. + """Represents a 'call' to a layer. + Connects input KerasTensors to output KerasTensors. """ + def __init__(self, layer, input_tensors, output_tensors): self.layer = layer self.input_tensors = input_tensors self.output_tensors = output_tensors - + # Register the node in the layer if not hasattr(layer, '_inbound_nodes'): layer._inbound_nodes = [] layer._inbound_nodes.append(self) - + # Link output tensors to this node if isinstance(output_tensors, list): for t in output_tensors: diff --git a/neutro/initializers/base.py b/neutro/initializers/base.py index c3d3eac..f6021e2 100644 --- a/neutro/initializers/base.py +++ b/neutro/initializers/base.py @@ -1,3 +1,6 @@ class Initializer: + """Base class for weight initializers.""" + def __call__(self, shape): + """Return an array of the given shape initialized per this strategy.""" raise NotImplementedError diff --git a/neutro/layers/__init__.py b/neutro/layers/__init__.py index 06d00ae..04b704c 100644 --- a/neutro/layers/__init__.py +++ b/neutro/layers/__init__.py @@ -4,12 +4,15 @@ from .core.dropout import Dropout from .core.flatten import Flatten from .core.activation import Activation, ReLU, Softmax, Sigmoid, Tanh +from .core.bitlinear import BitLinear +from .core.reparameterization import Reparameterization from .core.moe import MoELayer from .core.merging import Add, Concatenate, Multiply, Average, Maximum, Minimum from .convolutional.conv2d import Conv2D from .convolutional.conv1d import Conv1D from .pooling.maxpooling2d import MaxPooling2D from .pooling.global_pooling import GlobalAveragePooling2D, GlobalMaxPooling2D +from .pooling.upsampling2d import UpSampling2D from .recurrent.simple_rnn import SimpleRNN from .recurrent.lstm import LSTM from .recurrent.gru import GRU @@ -20,7 +23,10 @@ from .attention.paged_attention import PagedAttention, PagedKVCache from .normalization.layernorm import LayerNormalization from .normalization.batchnorm import BatchNormalization +from .normalization.rmsnorm import RMSNorm +from .normalization.groupnorm import GroupNormalization from .embedding.embedding import Embedding +from .embedding.time_embedding import TimeEmbedding from .embedding.token_position_embedding import TokenPositionEmbedding from .transformer.transformer_block import TransformerBlock from .transformer.bitnet_block import BitNetBlock diff --git a/neutro/layers/base.py b/neutro/layers/base.py index 1eacffd..b77c926 100644 --- a/neutro/layers/base.py +++ b/neutro/layers/base.py @@ -1,5 +1,6 @@ import numpy as np + class Layer: def __init__(self, name=None, **kwargs): self.name = name @@ -17,91 +18,96 @@ def build(self, input_shape): @property def sublayers(self): - """ - Returns all nested layers within this layer. + """Return all nested layers contained within this layer. + + Traverses instance attributes, including sublayers stored in lists, + tuples, or dictionaries, with cycle protection and de-duplication. """ layers = [] - for attr_name in dir(self): - if attr_name.startswith('_') or attr_name == 'sublayers': - continue - try: - attr = getattr(self, attr_name) - except AttributeError: - continue - + visited_ids = {id(self)} + stack = [] + for attr in vars(self).values(): + stack.append(attr) + + while stack: + attr = stack.pop() if isinstance(attr, Layer): - layers.append(attr) - elif isinstance(attr, list): - # Handle nested lists (like in MoELayer experts) - stack = [attr] - while stack: - curr = stack.pop() - for item in curr: - if isinstance(item, Layer): - layers.append(item) - elif isinstance(item, list): - stack.append(item) + if id(attr) not in visited_ids: + visited_ids.add(id(attr)) + layers.append(attr) + for nested in vars(attr).values(): + stack.append(nested) + elif isinstance(attr, (list, tuple, set)): + stack.extend(attr) + elif isinstance(attr, dict): + stack.extend(attr.values()) return layers def count_params(self): - """ - Counts the total number of parameters in this layer and its sublayers. - """ + """Count the total number of parameters in this layer and its sublayers.""" from neutro.autograd import Tensor as AutoTensor count = 0 for p in self.params.values(): - if isinstance(p, AutoTensor): - count += p.data.size - else: - count += p.size + count += p.data.size if isinstance(p, AutoTensor) else p.size for layer in self.sublayers: count += layer.count_params() return count def compute_output_shape(self, input_shape): - """ - Computes the output shape of the layer. - Should be overridden by subclasses. + """Compute the output shape of the layer. + + Should be overridden by subclasses that transform the input shape. """ if hasattr(self, 'output_shape') and self.output_shape is not None: return self.output_shape return input_shape - def backward(self, grad_output): - from neutro.autograd import Tensor as AT, GradientTape + def _collect_tensor_params(self): + """Gather all autograd Tensor parameters across the layer tree.""" + from neutro.autograd import Tensor as AT + tensor_params = {} + stack = [self] + while stack: + layer = stack.pop() + for param_name, param_value in layer.params.items(): + if isinstance(param_value, AT): + tensor_params[(id(layer), param_name)] = param_value + stack.extend(layer.sublayers) + return tensor_params + + def _captured_inputs(self): + """Build autograd Tensor input(s) from the last forward call.""" + from neutro.autograd import Tensor as AT inputs = getattr(self, '_last_inputs', None) if inputs is None: - return np.asarray(grad_output) - + return None if isinstance(inputs, list): - t_inputs = [AT(i) if not isinstance(i, AT) else i for i in inputs] - elif not isinstance(inputs, AT): - t_inputs = AT(np.asarray(inputs)) - else: - t_inputs = inputs + return [AT(i) if not isinstance(i, AT) else i for i in inputs] + return inputs if isinstance(inputs, AT) else AT(np.asarray(inputs)) + + def backward(self, grad_output): + from neutro.autograd import GradientTape, Tensor as AT + + t_inputs = self._captured_inputs() + if t_inputs is None: + return np.asarray(grad_output) fwd_args = getattr(self, '_last_args', ()) fwd_kwargs = getattr(self, '_last_kwargs', {}).copy() + # KV cache and layer id are runtime bookkeeping, not part of the + # differentiable computation. fwd_kwargs.pop('kv_cache', None) fwd_kwargs.pop('layer_id', None) - tensor_params = {} - stack_layers = [self] - while stack_layers: - l = stack_layers.pop() - for pn, pv in l.params.items(): - if isinstance(pv, AT): - tensor_params[(id(l), pn)] = pv - for sl in l.sublayers: - stack_layers.append(sl) - all_sources = list(tensor_params.values()) + + all_sources = list(self._collect_tensor_params().values()) if isinstance(t_inputs, list): all_sources.extend(t_inputs) else: all_sources.append(t_inputs) with GradientTape() as tape: - for s in all_sources: - tape.watch(s) + for source in all_sources: + tape.watch(source) output = self.forward(t_inputs, *fwd_args, **fwd_kwargs) g_t = AT(np.asarray(grad_output)) if isinstance(output, list): @@ -111,54 +117,53 @@ def backward(self, grad_output): tape.gradient(loss, all_sources) - stack_l = [self] - while stack_l: - l = stack_l.pop() - for pn, pv in l.params.items(): - from neutro.autograd import Tensor as AT2 - if isinstance(pv, AT2) and pv.grad is not None: - l.grads[pn] = pv.grad.copy() - for sl in l.sublayers: - stack_l.append(sl) + from neutro.autograd import Tensor as AT2 + stack_layers = [self] + while stack_layers: + layer = stack_layers.pop() + for param_name, param_value in layer.params.items(): + if isinstance(param_value, AT2) and param_value.grad is not None: + layer.grads[param_name] = param_value.grad.copy() + stack_layers.extend(layer.sublayers) if isinstance(t_inputs, list): return [t.grad.copy() if t.grad is not None else None for t in t_inputs] return t_inputs.grad.copy() if t_inputs.grad is not None else None - def __call__(self, inputs, *args, **kwargs): - from ..engine.node import KerasTensor, Node + def _is_symbolic_input(self, inputs): + from ..engine.node import KerasTensor - # Check if inputs are symbolic - is_symbolic = False if isinstance(inputs, KerasTensor): - is_symbolic = True - elif isinstance(inputs, list) and any(isinstance(i, KerasTensor) for i in inputs): - is_symbolic = True - - if is_symbolic: - # Symbolic call (Functional API) - if isinstance(inputs, list): - input_shapes = [i.shape for i in inputs] - else: - input_shapes = inputs.shape - - if not self.built: - self.build(input_shapes) - - output_shape = self.compute_output_shape(input_shapes) - - # Create output tensor(s) - if isinstance(output_shape, list): - output_tensors = [KerasTensor(shape=s) for s in output_shape] - else: - output_tensors = KerasTensor(shape=output_shape) - - # Create node - Node(self, input_tensors=inputs, output_tensors=output_tensors) - return output_tensors + return True + return (isinstance(inputs, list) + and any(isinstance(i, KerasTensor) for i in inputs)) + + def _symbolic_call(self, inputs): + """Execute the layer in graph-building (Functional API) mode.""" + from ..engine.node import KerasTensor, Node - # Eager call (Sequential or manual) + if isinstance(inputs, list): + input_shapes = [i.shape for i in inputs] + else: + input_shapes = inputs.shape + + if not self.built: + self.build(input_shapes) + + output_shape = self.compute_output_shape(input_shapes) + + if isinstance(output_shape, list): + output_tensors = [KerasTensor(shape=s) for s in output_shape] + else: + output_tensors = KerasTensor(shape=output_shape) + + Node(self, input_tensors=inputs, output_tensors=output_tensors) + return output_tensors + + def _eager_call(self, inputs, args, kwargs): + """Execute the layer in eager (Sequential/manual) mode.""" from neutro.autograd import as_tensor as _as_tensor + t_inputs = _as_tensor(inputs) if not self.built: if isinstance(t_inputs, list): @@ -172,6 +177,11 @@ def __call__(self, inputs, *args, **kwargs): self._last_output = output return output + def __call__(self, inputs, *args, **kwargs): + if self._is_symbolic_input(inputs): + return self._symbolic_call(inputs) + return self._eager_call(inputs, args, kwargs) + def get_params(self): return self.params diff --git a/neutro/layers/core/dense.py b/neutro/layers/core/dense.py index a5fece4..c1ec03e 100644 --- a/neutro/layers/core/dense.py +++ b/neutro/layers/core/dense.py @@ -6,7 +6,11 @@ class Dense(Layer): - def __init__(self, units, activation=None, use_bias=True, kernel_initializer='glorot_uniform', bias_initializer='zeros', **kwargs): + """A fully-connected (dense) layer.""" + + def __init__(self, units: int, activation=None, use_bias: bool = True, + kernel_initializer='glorot_uniform', bias_initializer='zeros', + **kwargs) -> None: super().__init__(**kwargs) self.units = units self.activation_name = activation @@ -15,7 +19,7 @@ def __init__(self, units, activation=None, use_bias=True, kernel_initializer='gl self.kernel_initializer = get_initializer(kernel_initializer) self.bias_initializer = get_initializer(bias_initializer) - def build(self, input_shape): + def build(self, input_shape) -> None: self.input_dim = input_shape[-1] self.params['W'] = Tensor(self.kernel_initializer((self.input_dim, self.units))) if self.use_bias: @@ -31,4 +35,4 @@ def forward(self, inputs, training=False): z = z + self.params['b'] if self.activation: return self.activation(z) - return z \ No newline at end of file + return z diff --git a/neutro/layers/core/merging.py b/neutro/layers/core/merging.py index f7320af..bf334f9 100644 --- a/neutro/layers/core/merging.py +++ b/neutro/layers/core/merging.py @@ -1,24 +1,24 @@ -import numpy as np from ..base import Layer from neutro.autograd import Tensor from neutro.autograd import ops as autograd_ops def _to_tensor(x): + """Coerce `x` into an autograd Tensor if it is not already one.""" return Tensor(x) if not isinstance(x, Tensor) else x -class Add(Layer): - def __init__(self, **kwargs): - super().__init__(**kwargs) +class _ReduceBase(Layer): + """Shared base for merging layers that reduce a list of tensors. + + All reduce-type layers (Add, Multiply, Average, Maximum, Minimum) + share identical build/output-shape logic and a forward pass that + folds a single binary operation over the input list. + """ def build(self, input_shape): - if isinstance(input_shape, list): - self.input_shape = input_shape - self.output_shape = input_shape[0] - else: - self.input_shape = input_shape - self.output_shape = input_shape + self.input_shape = input_shape + self.output_shape = self.compute_output_shape(input_shape) self.built = True def compute_output_shape(self, input_shape): @@ -26,15 +26,24 @@ def compute_output_shape(self, input_shape): return input_shape[0] return input_shape + def _combine(self, left, right): + """Combine two tensors. Must be implemented by subclasses.""" + raise NotImplementedError + def forward(self, inputs, training=False): if not isinstance(inputs, list): return inputs result = _to_tensor(inputs[0]) - for i in range(1, len(inputs)): - result = result + _to_tensor(inputs[i]) + for other in inputs[1:]: + result = self._combine(result, _to_tensor(other)) return result +class Add(_ReduceBase): + def _combine(self, left, right): + return left + right + + class Concatenate(Layer): def __init__(self, axis=-1, **kwargs): super().__init__(**kwargs) @@ -67,95 +76,27 @@ def forward(self, inputs, training=False): return autograd_ops.concatenate(tensors, axis=self.axis) -class Multiply(Layer): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def compute_output_shape(self, input_shape): - if isinstance(input_shape, list): - return input_shape[0] - return input_shape +class Multiply(_ReduceBase): + def _combine(self, left, right): + return left * right - def build(self, input_shape): - self.input_shape = input_shape - self.output_shape = self.compute_output_shape(input_shape) - self.built = True +class Average(_ReduceBase): def forward(self, inputs, training=False): - if not isinstance(inputs, list): - return inputs - result = _to_tensor(inputs[0]) - for i in range(1, len(inputs)): - result = result * _to_tensor(inputs[i]) + result = super().forward(inputs, training=training) + if isinstance(inputs, list) and len(inputs) > 1: + result = result / float(len(inputs)) return result + def _combine(self, left, right): + return left + right -class Average(Layer): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def compute_output_shape(self, input_shape): - if isinstance(input_shape, list): - return input_shape[0] - return input_shape - - def build(self, input_shape): - self.input_shape = input_shape - self.output_shape = self.compute_output_shape(input_shape) - self.built = True - def forward(self, inputs, training=False): - if not isinstance(inputs, list): - return inputs - result = _to_tensor(inputs[0]) - for i in range(1, len(inputs)): - result = result + _to_tensor(inputs[i]) - return result / float(len(inputs)) +class Maximum(_ReduceBase): + def _combine(self, left, right): + return autograd_ops.maximum(left, right) -class Maximum(Layer): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def compute_output_shape(self, input_shape): - if isinstance(input_shape, list): - return input_shape[0] - return input_shape - - def build(self, input_shape): - self.input_shape = input_shape - self.output_shape = self.compute_output_shape(input_shape) - self.built = True - - def forward(self, inputs, training=False): - if not isinstance(inputs, list): - return inputs - result = _to_tensor(inputs[0]) - for i in range(1, len(inputs)): - other = _to_tensor(inputs[i]) - result = autograd_ops.maximum(result, other) - return result - - -class Minimum(Layer): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def compute_output_shape(self, input_shape): - if isinstance(input_shape, list): - return input_shape[0] - return input_shape - - def build(self, input_shape): - self.input_shape = input_shape - self.output_shape = self.compute_output_shape(input_shape) - self.built = True - - def forward(self, inputs, training=False): - if not isinstance(inputs, list): - return inputs - result = _to_tensor(inputs[0]) - for i in range(1, len(inputs)): - other = _to_tensor(inputs[i]) - result = -autograd_ops.maximum(-result, -other) - return result \ No newline at end of file +class Minimum(_ReduceBase): + def _combine(self, left, right): + return -autograd_ops.maximum(-left, -right) \ No newline at end of file diff --git a/neutro/losses/base.py b/neutro/losses/base.py index b09f865..468d1d1 100644 --- a/neutro/losses/base.py +++ b/neutro/losses/base.py @@ -1,7 +1,15 @@ +import numpy as np + from neutro.autograd import Tensor + class Loss: + """Base class for all losses.""" + def __call__(self, y_true, y_pred): + """Compute the loss value between targets and predictions.""" raise NotImplementedError - def gradient(self, y_true, y_pred): + + def gradient(self, y_true, y_pred) -> np.ndarray: + """Compute the gradient of the loss w.r.t. the predictions.""" raise NotImplementedError diff --git a/neutro/losses/mse.py b/neutro/losses/mse.py index 8ea6517..276bf28 100644 --- a/neutro/losses/mse.py +++ b/neutro/losses/mse.py @@ -4,10 +4,12 @@ class MeanSquaredError(Loss): + """Mean squared error loss.""" + def __call__(self, y_true, y_pred): y_true = as_tensor(y_true) return ((y_pred - y_true) ** 2).mean() def gradient(self, y_true, y_pred): y_true = as_tensor(y_true) - return 2.0 * (y_pred - y_true) / y_true.data.size \ No newline at end of file + return 2.0 * (y_pred - y_true) / y_true.data.size diff --git a/neutro/metrics/base.py b/neutro/metrics/base.py index 30f6883..54a970d 100644 --- a/neutro/metrics/base.py +++ b/neutro/metrics/base.py @@ -1,5 +1,10 @@ class Metric: + """Base class for all metrics.""" + def __call__(self, y_true, y_pred): + """Compute the metric value between targets and predictions.""" raise NotImplementedError - def get_name(self): + + def get_name(self) -> str: + """Return the display name of the metric.""" raise NotImplementedError diff --git a/neutro/models/base_model.py b/neutro/models/base_model.py index 3866a96..264ace9 100644 --- a/neutro/models/base_model.py +++ b/neutro/models/base_model.py @@ -7,6 +7,7 @@ from ..layers.base import Layer + class Model(Layer): def __init__(self, inputs=None, outputs=None, name=None): super().__init__(name=name) @@ -15,26 +16,24 @@ def __init__(self, inputs=None, outputs=None, name=None): self.loss_fn = None self.metrics = [] self.stop_training = False - + self.inputs = inputs self.outputs = outputs - + if inputs is not None and outputs is not None: self._init_graph(inputs, outputs) def _init_graph(self, inputs, outputs): - """ - Traverses the graph from outputs to inputs to discover all layers and nodes. - """ + """Traverse the graph from outputs to inputs to discover layers and nodes.""" from ..engine.node import Node - + self._nodes_by_depth = [] self._layers = [] - + # Topological sort visited_nodes = set() nodes_ordered = [] - + def traverse(tensor): if hasattr(tensor, 'node') and tensor.node: node = tensor.node @@ -47,15 +46,15 @@ def traverse(tensor): else: traverse(node.input_tensors) nodes_ordered.append(node) - + if isinstance(outputs, list): for o in outputs: traverse(o) else: traverse(outputs) - + self._nodes_ordered = nodes_ordered - + # Collect all unique layers for node in nodes_ordered: if node.layer not in self.layers: @@ -74,7 +73,7 @@ def _get_all_layers(self, layers=None, visited=None): layers = self.layers if visited is None: visited = set() - + all_layers = [] for layer in layers: l_id = id(layer) @@ -91,7 +90,15 @@ def _get_all_layers(self, layers=None, visited=None): @staticmethod def _capture_layer_state(layer): """Recursively capture state of a layer and all its sublayers. - Returns dict: {id(sublayer): {attr_name: value, ...}}""" + + State is deep-copied so that later mutations during subsequent + forward passes (e.g. shared layers across branches) cannot corrupt + the captured snapshot. + + Returns dict: {id(sublayer): {attr_name: value, ...}} + """ + import copy + state = {} stack = [layer] visited = set() @@ -104,7 +111,7 @@ def _capture_layer_state(layer): sub = {} for k, v in l.__dict__.items(): if k not in Model._STATE_EXCLUDE: - sub[k] = v + sub[k] = copy.deepcopy(v) state[l_id] = sub for sl in l.sublayers: stack.append(sl) @@ -164,9 +171,22 @@ def _accumulate_layer_grads(layer, grads_accumulator): stack.append(sl) def fit(self, x, y=None, epochs=1, batch_size=32, verbose=1, validation_data=None, callbacks=None): + """Train the model using manual backward passes (layer.backward).""" + return self._fit_loop( + x=x, y=y, epochs=epochs, batch_size=batch_size, verbose=verbose, + validation_data=validation_data, callbacks=callbacks, use_autograd=False) + + def autograd_fit(self, x, y=None, epochs=1, batch_size=32, verbose=1, validation_data=None, callbacks=None): + """Train the model using the autograd engine (GradientTape).""" + return self._fit_loop( + x=x, y=y, epochs=epochs, batch_size=batch_size, verbose=verbose, + validation_data=validation_data, callbacks=callbacks, use_autograd=True) + + def _fit_loop(self, x, y=None, epochs=1, batch_size=32, verbose=1, + validation_data=None, callbacks=None, use_autograd=False): is_mimo_x = isinstance(x, list) is_mimo_y = isinstance(y, list) - + use_generator = False if not is_mimo_x and hasattr(x, '__iter__') and not isinstance(x, np.ndarray): use_generator = True @@ -176,99 +196,84 @@ def fit(self, x, y=None, epochs=1, batch_size=32, verbose=1, validation_data=Non n_samples = x[0].shape[0] else: n_samples = x.shape[0] - + history = History() history.set_model(self) - + all_callbacks = [history] + (callbacks or []) for cb in all_callbacks: cb.set_model(self) - + logs = {} - for cb in all_callbacks: cb.on_train_begin(logs) + for cb in all_callbacks: + cb.on_train_begin(logs) for epoch in range(epochs): - if self.stop_training: break - for cb in all_callbacks: cb.on_epoch_begin(epoch, logs) - + if self.stop_training: + break + for cb in all_callbacks: + cb.on_epoch_begin(epoch, logs) + epoch_loss = 0 epoch_metrics = {m.get_name(): 0 for m in self.metrics} - + if use_generator: num_batches = len(x) data_iter = iter(x) else: indices = np.arange(n_samples) np.random.shuffle(indices) - + if is_mimo_x: x_shuffled = [xi[indices] for xi in x] else: x_shuffled = x[indices] - + if is_mimo_y: y_shuffled = [yi[indices] for yi in y] else: y_shuffled = y[indices] - + num_batches = int(np.ceil(n_samples / batch_size)) - + total_seen = 0 if verbose == 1: - pbar = tqdm(total=num_batches, desc=f"Epoch {epoch+1}/{epochs}") - + pbar = tqdm(total=num_batches, desc=f"Epoch {epoch + 1}/{epochs}") + for i in range(num_batches): if use_generator: x_batch, y_batch = next(data_iter) else: start, end = i * batch_size, min((i + 1) * batch_size, n_samples) - + if is_mimo_x: x_batch = [xi[start:end] for xi in x_shuffled] else: x_batch = x_shuffled[start:end] - + if is_mimo_y: y_batch = [yi[start:end] for yi in y_shuffled] else: y_batch = y_shuffled[start:end] - + batch_size_actual = x_batch[0].shape[0] if is_mimo_x else len(x_batch) total_seen += batch_size_actual - - for cb in all_callbacks: cb.on_batch_begin(i, logs) - - # Forward - output = self.forward(x_batch, training=True) - - # Loss - sum across multiple outputs if applicable - is_mimo_out = isinstance(self.outputs, list) - if is_mimo_out: - batch_loss = sum(self.loss_fn(y_batch[j], output[j]) for j in range(len(self.outputs))) - else: - batch_loss = self.loss_fn(y_batch, output) + + for cb in all_callbacks: + cb.on_batch_begin(i, logs) + + batch_loss, output = self._train_on_batch(x_batch, y_batch, use_autograd) epoch_loss += batch_loss * batch_size_actual - + for m in self.metrics: try: m_val = m(y_batch, output) except (TypeError, ValueError): - m_val = m(y_batch[0], output[0]) if is_mimo_out else 0.0 + m_val = m(y_batch[0], output[0]) if isinstance(self.outputs, list) else 0.0 epoch_metrics[m.get_name()] += m_val * batch_size_actual - - # Backward - if is_mimo_out: - grads = [self.loss_fn.gradient(y_batch[j], output[j]) for j in range(len(self.outputs))] - self.backward(grads) - else: - grad = self.loss_fn.gradient(y_batch, output) - self.backward(grad) - - # Update - all_trainable_layers = self._get_all_layers() - self.optimizer.step(all_trainable_layers) - - for cb in all_callbacks: cb.on_batch_end(i, logs) + + for cb in all_callbacks: + cb.on_batch_end(i, logs) if verbose == 1: postfix = {'loss': f"{epoch_loss / total_seen:.4f}"} @@ -277,7 +282,7 @@ def fit(self, x, y=None, epochs=1, batch_size=32, verbose=1, validation_data=Non postfix[name] = f"{epoch_metrics[name] / total_seen:.4f}" pbar.set_postfix(postfix) pbar.update(1) - + if verbose == 1: pbar.close() @@ -285,24 +290,13 @@ def fit(self, x, y=None, epochs=1, batch_size=32, verbose=1, validation_data=Non 'loss': epoch_loss / total_seen, **{k: v / total_seen for k, v in epoch_metrics.items()} } - + if validation_data: - val_x, val_y = validation_data - val_output = self.predict(val_x) - is_mimo_val_out = isinstance(self.outputs, list) - if is_mimo_val_out: - logs['val_loss'] = sum(self.loss_fn(val_y[j], val_output[j]) for j in range(len(self.outputs))) - else: - logs['val_loss'] = self.loss_fn(val_y, val_output) - for m in self.metrics: - try: - m_val = m(val_y, val_output) - except (TypeError, ValueError): - m_val = m(val_y[0], val_output[0]) if is_mimo_val_out else 0.0 - logs[f'val_{m.get_name()}'] = m_val + self._update_validation_logs(logs, validation_data) + + for cb in all_callbacks: + cb.on_epoch_end(epoch, logs) - for cb in all_callbacks: cb.on_epoch_end(epoch, logs) - if verbose: if verbose == 1: if validation_data: @@ -312,7 +306,7 @@ def fit(self, x, y=None, epochs=1, batch_size=32, verbose=1, validation_data=Non val_msg += f" - val_{name}: {logs[f'val_{name}']:.4f}" print(val_msg) else: - msg = f"Epoch {epoch+1}/{epochs} - loss: {logs['loss']:.4f}" + msg = f"Epoch {epoch + 1}/{epochs} - loss: {logs['loss']:.4f}" for m in self.metrics: name = m.get_name() msg += f" - {name}: {logs[name]:.4f}" @@ -322,10 +316,72 @@ def fit(self, x, y=None, epochs=1, batch_size=32, verbose=1, validation_data=Non name = m.get_name() msg += f" - val_{name}: {logs[f'val_{name}']:.4f}" print(msg) - - for cb in all_callbacks: cb.on_train_end(logs) + + for cb in all_callbacks: + cb.on_train_end(logs) return history + def _train_on_batch(self, x_batch, y_batch, use_autograd): + """Run a single training step, returning (batch_loss, output).""" + if use_autograd: + from neutro.autograd import GradientTape, Tensor as AutoTensor + + with GradientTape() as tape: + for p in self.trainable_params: + tape.watch(p) + output = self.forward(x_batch, training=True) + if isinstance(self.outputs, list) and len(self.outputs) > 1: + if isinstance(y_batch, list): + loss = sum(self.loss_fn(y_batch[j], output[j]) for j in range(len(self.outputs))) + else: + loss = sum(self.loss_fn(y_batch, o) for o in output) + else: + loss = self.loss_fn(y_batch, output) + + batch_loss = float(loss.data) if isinstance(loss, AutoTensor) else float(loss) + + tape.gradient(loss, self.trainable_params) + + all_trainable_layers = self._get_all_layers() + self.optimizer.step(all_trainable_layers) + self._zero_all_grads() + else: + output = self.forward(x_batch, training=True) + + is_mimo_out = isinstance(self.outputs, list) + if is_mimo_out: + batch_loss = sum(self.loss_fn(y_batch[j], output[j]) for j in range(len(self.outputs))) + else: + batch_loss = self.loss_fn(y_batch, output) + + if is_mimo_out: + grads = [self.loss_fn.gradient(y_batch[j], output[j]) for j in range(len(self.outputs))] + self.backward(grads) + else: + grad = self.loss_fn.gradient(y_batch, output) + self.backward(grad) + + all_trainable_layers = self._get_all_layers() + self.optimizer.step(all_trainable_layers) + + return batch_loss, output + + def _update_validation_logs(self, logs, validation_data): + """Evaluate on validation data and write results into `logs`.""" + val_x, val_y = validation_data + val_output = self.predict(val_x) + is_mimo_val_out = isinstance(self.outputs, list) + if is_mimo_val_out: + logs['val_loss'] = sum(self.loss_fn(val_y[j], val_output[j]) for j in range(len(self.outputs))) + else: + logs['val_loss'] = self.loss_fn(val_y, val_output) + for m in self.metrics: + try: + m_val = m(val_y, val_output) + except (TypeError, ValueError): + m_val = m(val_y[0], val_output[0]) if is_mimo_val_out else 0.0 + logs[f'val_{m.get_name()}'] = m_val + @property def trainable_params(self): from neutro.autograd import Tensor as AutoTensor @@ -341,190 +397,37 @@ def _zero_all_grads(self): for p in self.trainable_params: p.zero_grad() - def autograd_fit(self, x, y=None, epochs=1, batch_size=32, verbose=1, validation_data=None, callbacks=None): - from neutro.autograd import GradientTape, Tensor as AutoTensor - from ..callbacks import History - - is_mimo_x = isinstance(x, list) - is_mimo_y = isinstance(y, list) - - if not is_mimo_x and hasattr(x, '__iter__') and not isinstance(x, np.ndarray): - n_samples = len(x) * x.batch_size if hasattr(x, 'batch_size') else len(x) - use_generator = True - else: - use_generator = False - if is_mimo_x: - n_samples = x[0].shape[0] - else: - n_samples = x.shape[0] - - history = History() - history.set_model(self) - all_callbacks = [history] + (callbacks or []) - for cb in all_callbacks: - cb.set_model(self) - - logs = {} - for cb in all_callbacks: cb.on_train_begin(logs) - - for epoch in range(epochs): - if self.stop_training: break - for cb in all_callbacks: cb.on_epoch_begin(epoch, logs) - - epoch_loss = 0 - epoch_metrics = {m.get_name(): 0 for m in self.metrics} - - if use_generator: - num_batches = len(x) - data_iter = iter(x) - else: - indices = np.arange(n_samples) - np.random.shuffle(indices) - if is_mimo_x: - x_shuffled = [xi[indices] for xi in x] - else: - x_shuffled = x[indices] - if is_mimo_y: - y_shuffled = [yi[indices] for yi in y] - else: - y_shuffled = y[indices] - num_batches = int(np.ceil(n_samples / batch_size)) - - total_seen = 0 - if verbose == 1: - pbar = tqdm(total=num_batches, desc=f"Epoch {epoch+1}/{epochs}") - - for i in range(num_batches): - if use_generator: - x_batch, y_batch = next(data_iter) - else: - start, end = i * batch_size, min((i + 1) * batch_size, n_samples) - if is_mimo_x: - x_batch = [xi[start:end] for xi in x_shuffled] - else: - x_batch = x_shuffled[start:end] - if is_mimo_y: - y_batch = [yi[start:end] for yi in y_shuffled] - else: - y_batch = y_shuffled[start:end] - - batch_size_actual = x_batch[0].shape[0] if is_mimo_x else len(x_batch) - total_seen += batch_size_actual - - for cb in all_callbacks: cb.on_batch_begin(i, logs) - - with GradientTape() as tape: - for p in self.trainable_params: - tape.watch(p) - output = self.forward(x_batch, training=True) - if isinstance(self.outputs, list) and len(self.outputs) > 1: - if is_mimo_y: - loss = sum(self.loss_fn(y_batch[j], output[j]) for j in range(len(self.outputs))) - else: - loss = sum(self.loss_fn(y_batch, o) for o in output) - else: - loss = self.loss_fn(y_batch, output) - - loss_data = float(loss.data) if isinstance(loss, AutoTensor) else float(loss) - epoch_loss += loss_data * batch_size_actual - - tape.gradient(loss, self.trainable_params) - - all_trainable_layers = self._get_all_layers() - self.optimizer.step(all_trainable_layers) - self._zero_all_grads() - - for m in self.metrics: - try: - m_val = m(y_batch, output) - except (TypeError, ValueError): - m_val = m(y_batch[0], output[0]) if isinstance(output, list) else 0.0 - epoch_metrics[m.get_name()] += m_val * batch_size_actual - - for cb in all_callbacks: cb.on_batch_end(i, logs) - - if verbose == 1: - postfix = {'loss': f"{epoch_loss / total_seen:.4f}"} - for m in self.metrics: - postfix[m.get_name()] = f"{epoch_metrics[m.get_name()] / total_seen:.4f}" - pbar.set_postfix(postfix) - pbar.update(1) - - if verbose == 1: - pbar.close() - - logs = { - 'loss': epoch_loss / total_seen, - **{k: v / total_seen for k, v in epoch_metrics.items()} - } - - if validation_data: - val_x, val_y = validation_data - val_output = self.predict(val_x) - if isinstance(self.outputs, list): - logs['val_loss'] = sum(self.loss_fn(val_y[j], val_output[j]) for j in range(len(self.outputs))) - else: - logs['val_loss'] = self.loss_fn(val_y, val_output) - for m in self.metrics: - try: - m_val = m(val_y, val_output) - except (TypeError, ValueError): - m_val = m(val_y[0], val_output[0]) if isinstance(self.outputs, list) else 0.0 - logs[f'val_{m.get_name()}'] = m_val - - for cb in all_callbacks: cb.on_epoch_end(epoch, logs) - - if verbose: - if verbose == 1: - if validation_data: - val_msg = f" - val_loss: {logs['val_loss']:.4f}" - for m in self.metrics: - val_msg += f" - val_{m.get_name()}: {logs[f'val_{m.get_name()}']:.4f}" - print(val_msg) - else: - msg = f"Epoch {epoch+1}/{epochs} - loss: {logs['loss']:.4f}" - for m in self.metrics: - msg += f" - {m.get_name()}: {logs[m.get_name()]:.4f}" - if validation_data: - msg += f" - val_loss: {logs['val_loss']:.4f}" - for m in self.metrics: - msg += f" - val_{m.get_name()}: {logs[f'val_{m.get_name()}']:.4f}" - print(msg) - - for cb in all_callbacks: cb.on_train_end(logs) - return history - def forward(self, inputs, training=False, kv_cache=None): if self.inputs is not None: # Functional API forward pass tensor_map = {} - + # Map input values if isinstance(self.inputs, list): for i, t in enumerate(self.inputs): tensor_map[id(t)] = inputs[i] else: tensor_map[id(self.inputs)] = inputs - + from ..layers.core.input_layer import InputLayer for node in self._nodes_ordered: if isinstance(node.layer, InputLayer): continue - + # Prepare inputs for this node if isinstance(node.input_tensors, list): node_inputs = [tensor_map.get(id(t)) for t in node.input_tensors] else: node_inputs = tensor_map.get(id(node.input_tensors)) - + output = node.layer.forward(node_inputs, training=training) from neutro.autograd import as_tensor node.layer._last_inputs = as_tensor(node_inputs) node.layer._last_kwargs = {'training': training} - + # Capture state AFTER forward so it captures the current call's data node.state = self._capture_layer_state(node.layer) - + # Store outputs if isinstance(node.output_tensors, list): for i, t in enumerate(node.output_tensors): @@ -532,7 +435,6 @@ def forward(self, inputs, training=False, kv_cache=None): else: tensor_map[id(node.output_tensors)] = output - # Return model outputs if isinstance(self.outputs, list): return [tensor_map[id(o)] for o in self.outputs] @@ -554,80 +456,76 @@ def forward(self, inputs, training=False, kv_cache=None): return inputs def generate(self, start_tokens, max_new_tokens, temperature=1.0): - """ - Autoregressive generation with KV Caching. - """ + """Autoregressive generation with KV Caching.""" from ..layers.attention.kv_cache import KVCache cache = KVCache() - + # Start with the full prompt curr_tokens = start_tokens generated = start_tokens - + for _ in range(max_new_tokens): # Forward pass logits = self.forward(curr_tokens, training=False, kv_cache=cache) - + # Get the logits for the last token only: (batch, seq, vocab) -> (batch, vocab) next_token_logits = logits[:, -1, :] / temperature - - # Simple sample or argmax - # For "naive" let's do argmax or simple categorical sample + + # Softmax probabilities probs = np.exp(next_token_logits - np.max(next_token_logits, axis=-1, keepdims=True)) probs /= np.sum(probs, axis=-1, keepdims=True) - + # Sample next token next_token = np.array([np.random.choice(len(p), p=p) for p in probs]) next_token = next_token.reshape(-1, 1) - - # Update inputs for next step - # With KV Cache, we only need to pass the LAST token + + # With KV Cache, only the last token needs to be passed forward curr_tokens = next_token generated = np.concatenate([generated, next_token], axis=1) - + return generated def backward(self, grad): if self.outputs is not None: # Functional API backward pass grad_map = {} - + # Map output gradients if isinstance(self.outputs, list): for i, t in enumerate(self.outputs): grad_map[id(t)] = grad[i] else: grad_map[id(self.outputs)] = grad - + # Initialize accumulators for shared layers layer_grads_accumulator = {} - + from ..layers.core.input_layer import InputLayer for node in reversed(self._nodes_ordered): if isinstance(node.layer, InputLayer): continue - + # Get gradients for this node's outputs if isinstance(node.output_tensors, list): node_grad_outputs = [grad_map.get(id(t)) for t in node.output_tensors] else: node_grad_outputs = grad_map.get(id(node.output_tensors)) - + if node_grad_outputs is None: continue - + # Restore state for this node recursively if hasattr(node, 'state'): self._restore_layer_state(node.layer, node.state) - + # Call layer.backward # Temporarily clear layer-tree grads to capture only this node call self._clear_layer_grads(node.layer) grad_inputs = node.layer.backward(node_grad_outputs) - + # Accumulate parameter gradients across the full layer tree self._accumulate_layer_grads(node.layer, layer_grads_accumulator) - + # Propagate gradients to inputs if isinstance(node.input_tensors, list): if not isinstance(grad_inputs, list): @@ -644,7 +542,7 @@ def backward(self, grad): grad_map[t_id] += grad_inputs else: grad_map[t_id] = grad_inputs - + # Return gradients for model inputs if isinstance(self.inputs, list): return [grad_map.get(id(i)) for i in self.inputs] @@ -661,11 +559,11 @@ def compute_output_shape(self, input_shape): if isinstance(self.outputs, list): return [o.shape for o in self.outputs] return self.outputs.shape - + # For Sequential models if not self.layers: return input_shape - + curr_shape = input_shape for layer in self.layers: curr_shape = layer.compute_output_shape(curr_shape) @@ -720,21 +618,19 @@ def evaluate(self, x, y): return results def summary(self): - """ - Prints a Keras-style summary of the model. - """ + """Print a Keras-style summary of the model.""" is_functional = self.inputs is not None - + print("-" * 85) if is_functional: print(f"{'Layer (type)':<25} {'Output Shape':<20} {'Param #':<10} {'Connected to':<25}") else: print(f"{'Layer (type)':<25} {'Output Shape':<20} {'Param #':<10}") print("=" * 85) - + total_params = 0 trainable_params = 0 - + curr_shape = None if self.layers and self.layers[0].input_shape: curr_shape = self.layers[0].input_shape @@ -742,7 +638,7 @@ def summary(self): for layer in self.layers: name = layer.name or layer.__class__.__name__ layer_type = layer.__class__.__name__ - + # Use layer's own input/output shapes if built if layer.built: try: @@ -757,12 +653,12 @@ def summary(self): output_shape = "multiple" else: output_shape = "unbuilt" - + params = layer.count_params() total_params += params if getattr(layer, 'trainable', True): trainable_params += params - + if is_functional: # Find which layers this layer is connected to via _inbound_nodes connected_to = [] @@ -776,12 +672,12 @@ def summary(self): connected_to.append(t.node.layer.name or t.node.layer.__class__.__name__) elif node.input_tensors and node.input_tensors.node: connected_to.append(node.input_tensors.node.layer.name or node.input_tensors.node.layer.__class__.__name__) - + connected_str = ", ".join(connected_to) if connected_to else "" print(f"{name + ' (' + layer_type + ')':<25} {str(output_shape):<20} {params:<10,} {connected_str:<25}") else: print(f"{name + ' (' + layer_type + ')':<25} {str(output_shape):<20} {params:<10,}") - + print("=" * 85) print(f"Total params: {total_params:,}") print(f"Trainable params: {trainable_params:,}") @@ -795,6 +691,7 @@ def save(self, filepath): def load(filepath): return joblib.load(filepath) + class Sequential(Model): def __init__(self, layers=None): super().__init__() @@ -811,9 +708,9 @@ def add(self, layer): # Convention: if rank is 1 (seq_len) or 2 (flat) or 3 (image), we might need batch # But it's safer to just check if the user provided it. # In LlamaTiny, input_shape=(seq_len,) is passed. - if len(shape) == 1: + if len(shape) == 1: shape = (None,) + shape - elif len(shape) == 3: # (h, w, c) + elif len(shape) == 3: # (h, w, c) shape = (None,) + shape layer.build(shape) else: @@ -822,5 +719,5 @@ def add(self, layer): if prev_layer.built: input_shape = prev_layer.compute_output_shape(prev_layer.input_shape) layer.build(input_shape) - + self.layers.append(layer) diff --git a/neutro/models/moe/__init__.py b/neutro/models/moe/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/neutro/models/vision/__init__.py b/neutro/models/vision/__init__.py index 3d86a35..5ef2b6b 100644 --- a/neutro/models/vision/__init__.py +++ b/neutro/models/vision/__init__.py @@ -1,2 +1,5 @@ from .alexnet import AlexNet from .vgg import VGG16, VGG19 +from .unet import UNet +from .vae import VAE +from .diffusion_model import DiffusionModel diff --git a/neutro/optimizers/adam.py b/neutro/optimizers/adam.py index 8fc1eb7..a6e8f22 100644 --- a/neutro/optimizers/adam.py +++ b/neutro/optimizers/adam.py @@ -2,13 +2,19 @@ from .base import Optimizer from neutro.autograd import Tensor +DEFAULT_LR = 0.001 +DEFAULT_BETA_1 = 0.9 +DEFAULT_BETA_2 = 0.999 +DEFAULT_EPSILON = 1e-7 + def _get_data(p): return p.data if isinstance(p, Tensor) else p class Adam(Optimizer): - def __init__(self, learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-7): + def __init__(self, learning_rate=DEFAULT_LR, beta_1=DEFAULT_BETA_1, + beta_2=DEFAULT_BETA_2, epsilon=DEFAULT_EPSILON): super().__init__(learning_rate) self.beta_1 = beta_1 self.beta_2 = beta_2 @@ -23,20 +29,25 @@ def step(self, layers): if not getattr(layer, 'trainable', True): continue for param_name, param_value in layer.params.items(): - grad = param_value.grad if isinstance(param_value, Tensor) and param_value.grad is not None else layer.grads.get(param_name) + grad = (param_value.grad + if isinstance(param_value, Tensor) and param_value.grad is not None + else layer.grads.get(param_name)) if grad is None: continue - param_data = _get_data(param_value) - key = (id(layer), param_name) + self._apply_step(layer, param_name, param_value, grad) + + def _apply_step(self, layer, param_name, param_value, grad): + param_data = _get_data(param_value) + key = (id(layer), param_name) - if key not in self.m: - self.m[key] = np.zeros_like(param_data) - self.v[key] = np.zeros_like(param_data) + if key not in self.m: + self.m[key] = np.zeros_like(param_data) + self.v[key] = np.zeros_like(param_data) - self.m[key] = self.beta_1 * self.m[key] + (1 - self.beta_1) * grad - self.v[key] = self.beta_2 * self.v[key] + (1 - self.beta_2) * (grad**2) + self.m[key] = self.beta_1 * self.m[key] + (1 - self.beta_1) * grad + self.v[key] = self.beta_2 * self.v[key] + (1 - self.beta_2) * (grad ** 2) - m_hat = self.m[key] / (1 - self.beta_1**self.t) - v_hat = self.v[key] / (1 - self.beta_2**self.t) + m_hat = self.m[key] / (1 - self.beta_1 ** self.t) + v_hat = self.v[key] / (1 - self.beta_2 ** self.t) - param_data -= self.learning_rate * m_hat / (np.sqrt(v_hat) + self.epsilon) \ No newline at end of file + param_data -= self.learning_rate * m_hat / (np.sqrt(v_hat) + self.epsilon) diff --git a/neutro/optimizers/adamw.py b/neutro/optimizers/adamw.py index 9e877b4..fb1bc2e 100644 --- a/neutro/optimizers/adamw.py +++ b/neutro/optimizers/adamw.py @@ -1,45 +1,22 @@ -import numpy as np -from .base import Optimizer -from neutro.autograd import Tensor +from .adam import Adam, _get_data +DEFAULT_WEIGHT_DECAY = 0.01 -def _get_data(p): - return p.data if isinstance(p, Tensor) else p +class AdamW(Adam): + """Adam with decoupled weight decay. -class AdamW(Optimizer): - def __init__(self, learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-7, weight_decay=0.01): - super().__init__(learning_rate) - self.beta_1 = beta_1 - self.beta_2 = beta_2 - self.epsilon = epsilon - self.weight_decay = weight_decay - self.m = {} - self.v = {} - self.t = 0 - - def step(self, layers): - self.t += 1 - for layer in layers: - if not getattr(layer, 'trainable', True): - continue - for param_name, param_value in layer.params.items(): - grad = param_value.grad if isinstance(param_value, Tensor) and param_value.grad is not None else layer.grads.get(param_name) - if grad is None: - continue - param_data = _get_data(param_value) - key = (id(layer), param_name) - - if key not in self.m: - self.m[key] = np.zeros_like(param_data) - self.v[key] = np.zeros_like(param_data) + The weight decay is applied directly to the parameter *before* the + moment updates, decoupled from the gradient (Loshchilov & Hutter, 2019). + """ - param_data -= self.learning_rate * self.weight_decay * param_data - - self.m[key] = self.beta_1 * self.m[key] + (1 - self.beta_1) * grad - self.v[key] = self.beta_2 * self.v[key] + (1 - self.beta_2) * (grad**2) - - m_hat = self.m[key] / (1 - self.beta_1**self.t) - v_hat = self.v[key] / (1 - self.beta_2**self.t) + def __init__(self, learning_rate=0.001, beta_1=0.9, beta_2=0.999, + epsilon=1e-7, weight_decay=DEFAULT_WEIGHT_DECAY): + super().__init__(learning_rate=learning_rate, beta_1=beta_1, + beta_2=beta_2, epsilon=epsilon) + self.weight_decay = weight_decay - param_data -= self.learning_rate * m_hat / (np.sqrt(v_hat) + self.epsilon) \ No newline at end of file + def _apply_step(self, layer, param_name, param_value, grad): + param_data = _get_data(param_value) + param_data -= self.learning_rate * self.weight_decay * param_data + super()._apply_step(layer, param_name, param_value, grad) diff --git a/neutro/optimizers/base.py b/neutro/optimizers/base.py index 457c3c5..20bc0e5 100644 --- a/neutro/optimizers/base.py +++ b/neutro/optimizers/base.py @@ -1,14 +1,20 @@ +from typing import List + + class Optimizer: - def __init__(self, learning_rate=0.001): + """Base class for all optimizers.""" + + def __init__(self, learning_rate: float = 0.001) -> None: self.learning_rate = learning_rate @property - def lr(self): + def lr(self) -> float: return self.learning_rate @lr.setter - def lr(self, value): + def lr(self, value: float) -> None: self.learning_rate = value - def step(self, layers): + def step(self, layers: List) -> None: + """Update trainable parameters given a list of layers.""" raise NotImplementedError diff --git a/neutro/optimizers/sgd.py b/neutro/optimizers/sgd.py index 9b13925..43844ef 100644 --- a/neutro/optimizers/sgd.py +++ b/neutro/optimizers/sgd.py @@ -2,24 +2,31 @@ from .base import Optimizer from neutro.autograd import Tensor +DEFAULT_LR = 0.01 + def _get_data(p): return p.data if isinstance(p, Tensor) else p class SGD(Optimizer): - def __init__(self, learning_rate=0.01, momentum=0.0, nesterov=False): + """Stochastic gradient descent with optional momentum and Nesterov.""" + + def __init__(self, learning_rate: float = DEFAULT_LR, momentum: float = 0.0, + nesterov: bool = False) -> None: super().__init__(learning_rate) self.momentum = momentum self.nesterov = nesterov self.velocities = {} - def step(self, layers): + def step(self, layers) -> None: for layer in layers: if not getattr(layer, 'trainable', True): continue for param_name, param_value in layer.params.items(): - grad = param_value.grad if isinstance(param_value, Tensor) and param_value.grad is not None else layer.grads.get(param_name) + grad = (param_value.grad + if isinstance(param_value, Tensor) and param_value.grad is not None + else layer.grads.get(param_name)) if grad is None: continue param_data = _get_data(param_value) @@ -28,12 +35,11 @@ def step(self, layers): if self.momentum > 0: if key not in self.velocities: self.velocities[key] = np.zeros_like(param_data) - v = self.velocities[key] - v_next = self.momentum * v - self.learning_rate * grad + v_next = self.momentum * self.velocities[key] - self.learning_rate * grad self.velocities[key] = v_next if self.nesterov: param_data += self.momentum * v_next - self.learning_rate * grad else: param_data += v_next else: - param_data -= self.learning_rate * grad \ No newline at end of file + param_data -= self.learning_rate * grad diff --git a/neutro/tokenizers/bpe.py b/neutro/tokenizers/bpe.py index 540e8ae..12b6331 100644 --- a/neutro/tokenizers/bpe.py +++ b/neutro/tokenizers/bpe.py @@ -1,17 +1,23 @@ -import regex as re import base64 +import json + +import regex as re + def get_stats(ids): + """Count the frequency of each adjacent pair of token ids.""" counts = {} for pair in zip(ids, ids[1:]): counts[pair] = counts.get(pair, 0) + 1 return counts + def merge(ids, pair, idx): + """Replace every occurrence of `pair` in `ids` with the new token `idx`.""" new_ids = [] i = 0 while i < len(ids): - if i < len(ids) - 1 and ids[i] == pair[0] and ids[i+1] == pair[1]: + if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]: new_ids.append(idx) i += 2 else: @@ -19,27 +25,24 @@ def merge(ids, pair, idx): i += 1 return new_ids + class BPETokenizer: - """ - Minimal BPE Tokenizer implementation. + """Minimal BPE Tokenizer implementation. + Educational and "intentionally naive". """ + def __init__(self): # byte -> id mapping self.encoder = {bytes([i]): i for i in range(256)} self.decoder = {i: bytes([i]) for i in range(256)} - self.special_tokens = {} # str -> int - self.inverse_special_tokens = {} # int -> str + self.special_tokens = {} # str -> int + self.inverse_special_tokens = {} # int -> str - def train(self, text, vocab_size, verbose=False): + def _train_loop(self, ids_list, vocab_size, verbose): + """Run the greedy BPE merge loop across a list of token-id sequences.""" assert vocab_size >= 256 num_merges = vocab_size - 256 - - # text to bytes - byte_chunks = [text.encode("utf-8")] - - # We need to maintain a list of token ids for each chunk - ids_list = [list(chunk) for chunk in byte_chunks] for i in range(num_merges): stats = {} @@ -48,20 +51,21 @@ def train(self, text, vocab_size, verbose=False): stats[pair] = stats.get(pair, 0) + 1 if not stats: break - + pair = max(stats, key=stats.get) idx = 256 + i - - # Update chunks + ids_list = [merge(ids, pair, idx) for ids in ids_list] - - # Update vocab + new_token_bytes = self.decoder[pair[0]] + self.decoder[pair[1]] self.encoder[new_token_bytes] = idx self.decoder[idx] = new_token_bytes - + if verbose: - print(f"merge {i+1}/{num_merges}: {pair} -> {idx} had {stats[pair]} occurrences") + print(f"merge {i + 1}/{num_merges}: {pair} -> {idx} had {stats[pair]} occurrences") + + def train(self, text, vocab_size, verbose=False): + self._train_loop([list(text.encode("utf-8"))], vocab_size, verbose) def encode(self, text): return self._encode_piece(text.encode("utf-8")) @@ -70,20 +74,20 @@ def _encode_piece(self, piece_bytes): # Optimized BPE encoding using a list of token ids # Initially, each byte is a token ids = list(piece_bytes) - + while len(ids) >= 2: # Find the pair that would be merged first (lowest rank in encoder) stats = get_stats(ids) pair = min(stats, key=lambda p: self.encoder.get(self.decoder[p[0]] + self.decoder[p[1]], float("inf"))) - + # If the best pair is not in our merges, we are done pair_bytes = self.decoder[pair[0]] + self.decoder[pair[1]] if pair_bytes not in self.encoder: break - + idx = self.encoder[pair_bytes] ids = merge(ids, pair, idx) - + return ids def decode(self, ids): @@ -97,10 +101,10 @@ def decode(self, ids): raise ValueError(f"Invalid token id: {idx}") return b"".join(parts).decode("utf-8", errors="replace") + class RegexTokenizer(BPETokenizer): - """ - BPE Tokenizer with Regex splitting (GPT style). - """ + """BPE Tokenizer with Regex splitting (GPT style).""" + # GPT-4 split pattern GPT4_SPLIT_PATTERN = r"""'(?i:[sdmtre lve])| \?|\p{L}+|\p{N}+|[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""" @@ -110,36 +114,9 @@ def __init__(self, pattern=None): self.compiled_pattern = re.compile(self.pattern) def train(self, text, vocab_size, verbose=False): - assert vocab_size >= 256 - num_merges = vocab_size - 256 - - # split text into chunks chunks = self.compiled_pattern.findall(text) - # convert chunks to byte ids ids_list = [list(chunk.encode("utf-8")) for chunk in chunks] - - for i in range(num_merges): - stats = {} - for ids in ids_list: - # get stats within each chunk - for pair in zip(ids, ids[1:]): - stats[pair] = stats.get(pair, 0) + 1 - - if not stats: - break - - pair = max(stats, key=stats.get) - idx = 256 + i - - # merge in all chunks - ids_list = [merge(ids, pair, idx) for ids in ids_list] - - new_token_bytes = self.decoder[pair[0]] + self.decoder[pair[1]] - self.encoder[new_token_bytes] = idx - self.decoder[idx] = new_token_bytes - - if verbose: - print(f"merge {i+1}/{num_merges}: {pair} -> {idx} had {stats[pair]} occurrences") + self._train_loop(ids_list, vocab_size, verbose) def register_special_tokens(self, special_tokens): # special_tokens: dict of str -> int @@ -156,10 +133,10 @@ def encode(self, text, allowed_special="none"): if not special_tokens: return self._encode_normal(text) - + special_pattern = "(" + "|".join(re.escape(k) for k in special_tokens.keys()) + ")" parts = re.split(special_pattern, text) - + ids = [] for part in parts: if part in special_tokens: @@ -176,7 +153,6 @@ def _encode_normal(self, text): return all_ids def save(self, prefix): - import json model = { "pattern": self.pattern, "encoder": {base64.b64encode(k).decode("ascii"): v for k, v in self.encoder.items()}, @@ -184,9 +160,8 @@ def save(self, prefix): } with open(f"{prefix}.json", "w") as f: json.dump(model, f) - + def load(self, prefix): - import json with open(f"{prefix}.json", "r") as f: model = json.load(f) self.pattern = model["pattern"] @@ -199,4 +174,3 @@ def load(self, prefix): @property def vocab_size(self): return len(self.encoder) + len(self.special_tokens) - From fe0584205c27220650fd1c85ebebb6e2daa2976c Mon Sep 17 00:00:00 2001 From: "opencode-agent[bot]" Date: Sat, 1 Aug 2026 10:22:07 +0000 Subject: [PATCH 2/3] Pythonic refactor; 379 tests pass. Co-authored-by: sourcepirate --- neutro/callbacks/base.py | 6 +- neutro/data.py | 5 +- neutro/layers/base.py | 18 ++-- neutro/layers/core/dense.py | 2 +- neutro/layers/core/merging.py | 7 +- neutro/models/base_model.py | 196 ++++++++++++++-------------------- neutro/tokenizers/bpe.py | 15 ++- 7 files changed, 96 insertions(+), 153 deletions(-) diff --git a/neutro/callbacks/base.py b/neutro/callbacks/base.py index 67dd60c..1f8d0ff 100644 --- a/neutro/callbacks/base.py +++ b/neutro/callbacks/base.py @@ -51,11 +51,7 @@ def _is_improvement(self, current: float) -> bool: if self.mode == 'max': return current > self.best # auto: infer direction from the monitored metric's name - if 'acc' in self.monitor: - return current > self.best - if 'loss' in self.monitor: - return current < self.best - return current < self.best + return current > self.best if 'acc' in self.monitor else current < self.best def _init_best(self) -> float: maximize = self.mode == 'max' or (self.mode == 'auto' and 'acc' in self.monitor) diff --git a/neutro/data.py b/neutro/data.py index 43b6b2d..4538d6b 100644 --- a/neutro/data.py +++ b/neutro/data.py @@ -34,10 +34,7 @@ def __getitem__(self, index): batch_x, batch_y = self.x[indices], self.y[indices] if self.augmenter: - augmented_x = np.zeros_like(batch_x) - for i in range(len(batch_x)): - augmented_x[i] = self.augmenter.apply_transform(batch_x[i]) - batch_x = augmented_x + batch_x = np.stack([self.augmenter.apply_transform(xi) for xi in batch_x]) return batch_x, batch_y diff --git a/neutro/layers/base.py b/neutro/layers/base.py index b77c926..781c873 100644 --- a/neutro/layers/base.py +++ b/neutro/layers/base.py @@ -46,21 +46,16 @@ def sublayers(self): def count_params(self): """Count the total number of parameters in this layer and its sublayers.""" from neutro.autograd import Tensor as AutoTensor - count = 0 - for p in self.params.values(): - count += p.data.size if isinstance(p, AutoTensor) else p.size - for layer in self.sublayers: - count += layer.count_params() - return count + own = sum(p.data.size if isinstance(p, AutoTensor) else p.size + for p in self.params.values()) + return own + sum(layer.count_params() for layer in self.sublayers) def compute_output_shape(self, input_shape): """Compute the output shape of the layer. Should be overridden by subclasses that transform the input shape. """ - if hasattr(self, 'output_shape') and self.output_shape is not None: - return self.output_shape - return input_shape + return self.output_shape if self.output_shape is not None else input_shape def _collect_tensor_params(self): """Gather all autograd Tensor parameters across the layer tree.""" @@ -133,9 +128,8 @@ def backward(self, grad_output): def _is_symbolic_input(self, inputs): from ..engine.node import KerasTensor - if isinstance(inputs, KerasTensor): - return True - return (isinstance(inputs, list) + return (isinstance(inputs, KerasTensor) + or isinstance(inputs, list) and any(isinstance(i, KerasTensor) for i in inputs)) def _symbolic_call(self, inputs): diff --git a/neutro/layers/core/dense.py b/neutro/layers/core/dense.py index c1ec03e..b35779a 100644 --- a/neutro/layers/core/dense.py +++ b/neutro/layers/core/dense.py @@ -27,7 +27,7 @@ def build(self, input_shape) -> None: super().build(input_shape) def compute_output_shape(self, input_shape): - return tuple(list(input_shape)[:-1] + [self.units]) + return (*input_shape[:-1], self.units) def forward(self, inputs, training=False): z = inputs @ self.params['W'] diff --git a/neutro/layers/core/merging.py b/neutro/layers/core/merging.py index bf334f9..097408a 100644 --- a/neutro/layers/core/merging.py +++ b/neutro/layers/core/merging.py @@ -1,3 +1,5 @@ +from functools import reduce + from ..base import Layer from neutro.autograd import Tensor from neutro.autograd import ops as autograd_ops @@ -33,10 +35,7 @@ def _combine(self, left, right): def forward(self, inputs, training=False): if not isinstance(inputs, list): return inputs - result = _to_tensor(inputs[0]) - for other in inputs[1:]: - result = self._combine(result, _to_tensor(other)) - return result + return reduce(self._combine, map(_to_tensor, inputs)) class Add(_ReduceBase): diff --git a/neutro/models/base_model.py b/neutro/models/base_model.py index 264ace9..3e12ecf 100644 --- a/neutro/models/base_model.py +++ b/neutro/models/base_model.py @@ -1,3 +1,6 @@ +import copy +import inspect + import numpy as np import joblib from tqdm import tqdm @@ -8,6 +11,19 @@ from ..layers.base import Layer +def _iter_layer_tree(root): + """Yield each layer in `root`'s layer tree exactly once, cycle-safe.""" + stack = [root] + seen = set() + while stack: + layer = stack.pop() + if id(layer) in seen: + continue + seen.add(id(layer)) + yield layer + stack.extend(reversed(layer.sublayers)) + + class Model(Layer): def __init__(self, inputs=None, outputs=None, name=None): super().__init__(name=name) @@ -69,20 +85,18 @@ def compile(self, optimizer, loss, metrics=None): self.metrics = [metrics_module.get(m) for m in (metrics or [])] def _get_all_layers(self, layers=None, visited=None): - if layers is None: - layers = self.layers - if visited is None: - visited = set() - - all_layers = [] - for layer in layers: - l_id = id(layer) - if l_id not in visited: - visited.add(l_id) - all_layers.append(layer) - if hasattr(layer, 'sublayers'): - all_layers.extend(self._get_all_layers(layer.sublayers, visited)) - return all_layers + roots = self.layers if layers is None else layers + seen = set() if visited is None else visited + stack = list(reversed(roots)) + result = [] + while stack: + layer = stack.pop() + if id(layer) in seen: + continue + seen.add(id(layer)) + result.append(layer) + stack.extend(reversed(layer.sublayers)) + return result _STATE_EXCLUDE = {'params', 'grads', 'built', 'input_shape', 'output_shape', 'name', '_inbound_nodes', 'trainable'} @@ -97,69 +111,28 @@ def _capture_layer_state(layer): Returns dict: {id(sublayer): {attr_name: value, ...}} """ - import copy - - state = {} - stack = [layer] - visited = set() - while stack: - l = stack.pop() - l_id = id(l) - if l_id in visited: - continue - visited.add(l_id) - sub = {} - for k, v in l.__dict__.items(): - if k not in Model._STATE_EXCLUDE: - sub[k] = copy.deepcopy(v) - state[l_id] = sub - for sl in l.sublayers: - stack.append(sl) - return state + return { + id(l): {k: copy.deepcopy(v) for k, v in l.__dict__.items() + if k not in Model._STATE_EXCLUDE} + for l in _iter_layer_tree(layer) + } @staticmethod def _restore_layer_state(layer, state): """Restore state captured by _capture_layer_state onto layer tree.""" - stack = [layer] - visited = set() - while stack: - l = stack.pop() - l_id = id(l) - if l_id in visited: - continue - visited.add(l_id) - if l_id in state: - for k, v in state[l_id].items(): - setattr(l, k, v) - for sl in l.sublayers: - stack.append(sl) + for l in _iter_layer_tree(layer): + for k, v in state.get(id(l), {}).items(): + setattr(l, k, v) @staticmethod def _clear_layer_grads(layer): - stack = [layer] - visited = set() - while stack: - l = stack.pop() - l_id = id(l) - if l_id in visited: - continue - visited.add(l_id) + for l in _iter_layer_tree(layer): l.grads = {} - for sl in l.sublayers: - stack.append(sl) @staticmethod def _accumulate_layer_grads(layer, grads_accumulator): - stack = [layer] - visited = set() - while stack: - l = stack.pop() - l_id = id(l) - if l_id in visited: - continue - visited.add(l_id) - - layer_acc = grads_accumulator.setdefault(l_id, {}) + for l in _iter_layer_tree(layer): + layer_acc = grads_accumulator.setdefault(id(l), {}) for k, v in l.grads.items(): if k in layer_acc: layer_acc[k] += v @@ -167,8 +140,26 @@ def _accumulate_layer_grads(layer, grads_accumulator): layer_acc[k] = np.array(v, copy=True) l.grads = layer_acc - for sl in l.sublayers: - stack.append(sl) + @staticmethod + def _is_mimo(outputs): + return isinstance(outputs, list) + + def _compute_loss(self, y, output): + if self._is_mimo(self.outputs): + return sum(self.loss_fn(y[j], output[j]) for j in range(len(self.outputs))) + return self.loss_fn(y, output) + + @staticmethod + def _eval_metric(metric, y, output, is_mimo_out): + try: + return metric(y, output) + except (TypeError, ValueError): + return metric(y[0], output[0]) if is_mimo_out else 0.0 + + def _metric_summary(self, logs, prefix=''): + return ''.join( + f" - {prefix}{m.get_name()}: {logs[f'{prefix}{m.get_name()}']:.4f}" + for m in self.metrics) def fit(self, x, y=None, epochs=1, batch_size=32, verbose=1, validation_data=None, callbacks=None): """Train the model using manual backward passes (layer.backward).""" @@ -266,10 +257,7 @@ def _fit_loop(self, x, y=None, epochs=1, batch_size=32, verbose=1, epoch_loss += batch_loss * batch_size_actual for m in self.metrics: - try: - m_val = m(y_batch, output) - except (TypeError, ValueError): - m_val = m(y_batch[0], output[0]) if isinstance(self.outputs, list) else 0.0 + m_val = self._eval_metric(m, y_batch, output, self._is_mimo(self.outputs)) epoch_metrics[m.get_name()] += m_val * batch_size_actual for cb in all_callbacks: @@ -300,21 +288,13 @@ def _fit_loop(self, x, y=None, epochs=1, batch_size=32, verbose=1, if verbose: if verbose == 1: if validation_data: - val_msg = f" - val_loss: {logs['val_loss']:.4f}" - for m in self.metrics: - name = m.get_name() - val_msg += f" - val_{name}: {logs[f'val_{name}']:.4f}" - print(val_msg) + print(f" - val_loss: {logs['val_loss']:.4f}" + + self._metric_summary(logs, 'val_')) else: - msg = f"Epoch {epoch + 1}/{epochs} - loss: {logs['loss']:.4f}" - for m in self.metrics: - name = m.get_name() - msg += f" - {name}: {logs[name]:.4f}" + msg = (f"Epoch {epoch + 1}/{epochs} - loss: {logs['loss']:.4f}" + + self._metric_summary(logs)) if validation_data: - msg += f" - val_loss: {logs['val_loss']:.4f}" - for m in self.metrics: - name = m.get_name() - msg += f" - val_{name}: {logs[f'val_{name}']:.4f}" + msg += f" - val_loss: {logs['val_loss']:.4f}" + self._metric_summary(logs, 'val_') print(msg) for cb in all_callbacks: @@ -348,11 +328,8 @@ def _train_on_batch(self, x_batch, y_batch, use_autograd): else: output = self.forward(x_batch, training=True) - is_mimo_out = isinstance(self.outputs, list) - if is_mimo_out: - batch_loss = sum(self.loss_fn(y_batch[j], output[j]) for j in range(len(self.outputs))) - else: - batch_loss = self.loss_fn(y_batch, output) + is_mimo_out = self._is_mimo(self.outputs) + batch_loss = self._compute_loss(y_batch, output) if is_mimo_out: grads = [self.loss_fn.gradient(y_batch[j], output[j]) for j in range(len(self.outputs))] @@ -370,16 +347,10 @@ def _update_validation_logs(self, logs, validation_data): """Evaluate on validation data and write results into `logs`.""" val_x, val_y = validation_data val_output = self.predict(val_x) - is_mimo_val_out = isinstance(self.outputs, list) - if is_mimo_val_out: - logs['val_loss'] = sum(self.loss_fn(val_y[j], val_output[j]) for j in range(len(self.outputs))) - else: - logs['val_loss'] = self.loss_fn(val_y, val_output) + is_mimo_val_out = self._is_mimo(self.outputs) + logs['val_loss'] = self._compute_loss(val_y, val_output) for m in self.metrics: - try: - m_val = m(val_y, val_output) - except (TypeError, ValueError): - m_val = m(val_y[0], val_output[0]) if is_mimo_val_out else 0.0 + m_val = self._eval_metric(m, val_y, val_output, is_mimo_val_out) logs[f'val_{m.get_name()}'] = m_val @property @@ -443,14 +414,9 @@ def forward(self, inputs, training=False, kv_cache=None): # Sequential or Subclassed forward pass for i, layer in enumerate(self.layers): - if kv_cache is not None and hasattr(layer, 'forward'): - # Check if layer accepts kv_cache (Attention or Blocks) - import inspect - sig = inspect.signature(layer.forward) - if 'kv_cache' in sig.parameters: - inputs = layer(inputs, training=training, kv_cache=kv_cache, layer_id=i) - else: - inputs = layer(inputs, training=training) + if kv_cache is not None and 'kv_cache' in inspect.signature(layer.forward).parameters: + # Layer accepts kv_cache (Attention or Blocks) + inputs = layer(inputs, training=training, kv_cache=kv_cache, layer_id=i) else: inputs = layer(inputs, training=training) return inputs @@ -602,18 +568,12 @@ def _to_scalar(v): return float(v) def evaluate(self, x, y): - is_mimo_out = isinstance(self.outputs, list) + is_mimo_out = self._is_mimo(self.outputs) output = self.predict(x) - if is_mimo_out: - loss = sum(self.loss_fn(y[j], output[j]) for j in range(len(self.outputs))) - else: - loss = self.loss_fn(y, output) + loss = self._compute_loss(y, output) results = {'loss': self._to_scalar(loss)} for m in self.metrics: - try: - m_val = m(y, output) - except (TypeError, ValueError): - m_val = m(y[0], output[0]) if is_mimo_out else 0.0 + m_val = self._eval_metric(m, y, output, is_mimo_out) results[m.get_name()] = self._to_scalar(m_val) if not isinstance(m_val, float) else m_val return results @@ -674,9 +634,9 @@ def summary(self): connected_to.append(node.input_tensors.node.layer.name or node.input_tensors.node.layer.__class__.__name__) connected_str = ", ".join(connected_to) if connected_to else "" - print(f"{name + ' (' + layer_type + ')':<25} {str(output_shape):<20} {params:<10,} {connected_str:<25}") + print(f"{f'{name} ({layer_type})':<25} {output_shape!s:<20} {params:<10,} {connected_str:<25}") else: - print(f"{name + ' (' + layer_type + ')':<25} {str(output_shape):<20} {params:<10,}") + print(f"{f'{name} ({layer_type})':<25} {output_shape!s:<20} {params:<10,}") print("=" * 85) print(f"Total params: {total_params:,}") diff --git a/neutro/tokenizers/bpe.py b/neutro/tokenizers/bpe.py index 12b6331..1fffcaa 100644 --- a/neutro/tokenizers/bpe.py +++ b/neutro/tokenizers/bpe.py @@ -1,15 +1,14 @@ import base64 +import itertools import json +from collections import Counter import regex as re def get_stats(ids): """Count the frequency of each adjacent pair of token ids.""" - counts = {} - for pair in zip(ids, ids[1:]): - counts[pair] = counts.get(pair, 0) + 1 - return counts + return Counter(itertools.pairwise(ids)) def merge(ids, pair, idx): @@ -45,10 +44,8 @@ def _train_loop(self, ids_list, vocab_size, verbose): num_merges = vocab_size - 256 for i in range(num_merges): - stats = {} - for ids in ids_list: - for pair in zip(ids, ids[1:]): - stats[pair] = stats.get(pair, 0) + 1 + stats = Counter(itertools.chain.from_iterable( + itertools.pairwise(ids) for ids in ids_list)) if not stats: break @@ -110,7 +107,7 @@ class RegexTokenizer(BPETokenizer): def __init__(self, pattern=None): super().__init__() - self.pattern = pattern if pattern else self.GPT4_SPLIT_PATTERN + self.pattern = pattern or self.GPT4_SPLIT_PATTERN self.compiled_pattern = re.compile(self.pattern) def train(self, text, vocab_size, verbose=False): From 3db2a02df09e3342d809655410a4bbb3ea3dc5a2 Mon Sep 17 00:00:00 2001 From: "opencode-agent[bot]" Date: Sat, 1 Aug 2026 11:25:06 +0000 Subject: [PATCH 3/3] Autograd engine pythonified; 379 tests pass. Co-authored-by: sourcepirate --- neutro/autograd/custom_ops.py | 5 +- neutro/autograd/function.py | 10 +- neutro/autograd/ops.py | 732 ++++++++++++++++++---------------- neutro/autograd/tape.py | 26 +- neutro/autograd/tensor.py | 11 +- 5 files changed, 419 insertions(+), 365 deletions(-) diff --git a/neutro/autograd/custom_ops.py b/neutro/autograd/custom_ops.py index 94fca83..0e372c8 100644 --- a/neutro/autograd/custom_ops.py +++ b/neutro/autograd/custom_ops.py @@ -1,8 +1,5 @@ import numpy as np -from .tensor import Tensor -from .function import Function, _Ctx -from .utils import broadcast_backward -from .tape import get_active_tape +from .function import Function class Gather(Function): diff --git a/neutro/autograd/function.py b/neutro/autograd/function.py index 55d712e..201c991 100644 --- a/neutro/autograd/function.py +++ b/neutro/autograd/function.py @@ -34,7 +34,7 @@ def apply(cls, *args, **kwargs): result = Tensor(output_data) tape = get_active_tape() - if tape and any(id(t) in tape._watched for t in tensor_args): + if tape and any(t in tape._watched for t in tensor_args): for name in ctx.saved_data: val = ctx.saved_data[name] if isinstance(val, Tensor): @@ -47,13 +47,13 @@ def bw(g): grad_inputs = cls.backward(ctx_copy, g) if not isinstance(grad_inputs, (list, tuple)): grad_inputs = [grad_inputs] - result = [] + out_grads = [] for t, gi in zip(tensor_args, grad_inputs): if gi is not None: - result.append(broadcast_backward(gi, t.shape)) + out_grads.append(broadcast_backward(gi, t.shape)) else: - result.append(None) - return result + out_grads.append(None) + return out_grads tape._record_op(tensor_args, result, bw, cls.__name__) diff --git a/neutro/autograd/ops.py b/neutro/autograd/ops.py index 2df9a91..b33c712 100644 --- a/neutro/autograd/ops.py +++ b/neutro/autograd/ops.py @@ -1,7 +1,6 @@ import numpy as np from .tensor import Tensor -from .tape import get_active_tape -from .utils import broadcast_backward +from .function import Function def _ensure_tensor(x): @@ -10,423 +9,476 @@ def _ensure_tensor(x): return Tensor(x) +class _Add(Function): + @staticmethod + def forward(ctx, a, b): + ctx.save_for_backward(a=a, b=b) + return a + b + + @staticmethod + def backward(ctx, g): + a, b = ctx.saved_data['a'], ctx.saved_data['b'] + return g, g + + +class _Sub(Function): + @staticmethod + def forward(ctx, a, b): + ctx.save_for_backward(a=a, b=b) + return a - b + + @staticmethod + def backward(ctx, g): + return g, -g + + +class _Mul(Function): + @staticmethod + def forward(ctx, a, b): + ctx.save_for_backward(a=a, b=b) + return a * b + + @staticmethod + def backward(ctx, g): + a, b = ctx.saved_data['a'], ctx.saved_data['b'] + return g * b, g * a + + +class _Div(Function): + @staticmethod + def forward(ctx, a, b): + ctx.save_for_backward(a=a, b=b) + return a / b + + @staticmethod + def backward(ctx, g): + a, b = ctx.saved_data['a'], ctx.saved_data['b'] + return g / b, -g * a / (b ** 2) + + +class _Neg(Function): + @staticmethod + def forward(ctx, x): + return -x + + @staticmethod + def backward(ctx, g): + return -g + + +class _Pow(Function): + @staticmethod + def forward(ctx, x, k): + ctx.save_for_backward(x=x, k=k) + return x ** k + + @staticmethod + def backward(ctx, g): + x, k = ctx.saved_data['x'], ctx.saved_data['k'] + return k * (x ** (k - 1)) * g + + +class _Matmul(Function): + @staticmethod + def forward(ctx, a, b): + ctx.save_for_backward(a=a, b=b) + return a @ b + + @staticmethod + def backward(ctx, g): + a, b = ctx.saved_data['a'], ctx.saved_data['b'] + ga = g @ b.T if g.ndim == 2 and b.ndim == 2 else g @ np.swapaxes(b, -1, -2) + gb = a.T @ g if g.ndim == 2 and a.ndim == 2 else np.swapaxes(a, -1, -2) @ g + return ga, gb + + +class _Sum(Function): + @staticmethod + def forward(ctx, x, axis, keepdims): + ctx.save_for_backward(x_shape=x.shape, axis=axis, keepdims=keepdims) + return np.sum(x, axis=axis, keepdims=keepdims) + + @staticmethod + def backward(ctx, g): + x_shape, axis, keepdims = ctx.saved_data['x_shape'], ctx.saved_data['axis'], ctx.saved_data['keepdims'] + if axis is None: + gx = g * np.ones(x_shape, dtype=float) + else: + gx = np.expand_dims(g, axis=axis) if not keepdims else g + gx = gx * np.ones(x_shape, dtype=float) + return gx.reshape(x_shape) + + +class _Mean(Function): + @staticmethod + def forward(ctx, x, axis, keepdims): + if axis is None: + n = x.size + else: + axes = (axis,) if isinstance(axis, int) else axis + n = int(np.prod([x.shape[d] for d in axes])) + ctx.save_for_backward(x_shape=x.shape, axis=axis, keepdims=keepdims, n=n) + return np.mean(x, axis=axis, keepdims=keepdims) + + @staticmethod + def backward(ctx, g): + x_shape, axis, keepdims, n = ctx.saved_data['x_shape'], ctx.saved_data['axis'], ctx.saved_data['keepdims'], ctx.saved_data['n'] + if axis is None: + gx = g * np.ones(x_shape, dtype=float) / n + else: + gx = np.expand_dims(g, axis=axis) if not keepdims else g + gx = gx * np.ones(x_shape, dtype=float) / n + return gx.reshape(x_shape) + + +class _Transpose(Function): + @staticmethod + def forward(ctx, x, axes): + if axes is None: + axes = tuple(range(x.ndim - 1, -1, -1)) + ctx.save_for_backward(inv_axes=tuple(np.argsort(axes))) + return np.transpose(x, axes) + + @staticmethod + def backward(ctx, g): + return np.transpose(g, ctx.saved_data['inv_axes']) + + +class _Reshape(Function): + @staticmethod + def forward(ctx, x, shape): + ctx.save_for_backward(x_shape=x.shape) + return x.reshape(shape) + + @staticmethod + def backward(ctx, g): + return g.reshape(ctx.saved_data['x_shape']) + + +class _Slice(Function): + @staticmethod + def forward(ctx, x, idx): + ctx.save_for_backward(x=x, idx=idx) + return x[idx] + + @staticmethod + def backward(ctx, g): + full = np.zeros_like(ctx.saved_data['x']) + full[ctx.saved_data['idx']] += g + return full + + +class _Relu(Function): + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x=x) + return np.maximum(0, x) + + @staticmethod + def backward(ctx, g): + return g * (ctx.saved_data['x'] > 0).astype(float) + + +class _Sigmoid(Function): + @staticmethod + def forward(ctx, x): + s = 1.0 / (1.0 + np.exp(-np.clip(x, -500, 500))) + ctx.save_for_backward(s=s.copy()) + return s.copy() + + @staticmethod + def backward(ctx, g): + s = ctx.saved_data['s'] + return g * s * (1.0 - s) + + +class _Tanh(Function): + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x=x) + return np.tanh(x) + + @staticmethod + def backward(ctx, g): + t = np.tanh(ctx.saved_data['x']) + return g * (1.0 - t ** 2) + + +class _Silu(Function): + @staticmethod + def forward(ctx, x): + s = 1.0 / (1.0 + np.exp(-np.clip(x, -500, 500))) + ctx.save_for_backward(x=x, s=s.copy()) + return x * s + + @staticmethod + def backward(ctx, g): + x, s = ctx.saved_data['x'], ctx.saved_data['s'] + return g * (s + x * s * (1.0 - s)) + + +class _Softmax(Function): + @staticmethod + def forward(ctx, x, axis): + x_max = np.max(x, axis=axis, keepdims=True) + exps = np.exp(x - x_max) + s = exps / np.sum(exps, axis=axis, keepdims=True) + ctx.save_for_backward(s=s.copy(), axis=axis) + return s.copy() + + @staticmethod + def backward(ctx, g): + s, axis = ctx.saved_data['s'], ctx.saved_data['axis'] + dot = np.sum(s * g, axis=axis, keepdims=True) + return s * (g - dot) + + +class _Sqrt(Function): + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x=x) + return np.sqrt(x) + + @staticmethod + def backward(ctx, g): + return g / (2.0 * np.sqrt(ctx.saved_data['x']) + 1e-15) + + +class _Exp(Function): + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x=x) + return np.exp(x) + + @staticmethod + def backward(ctx, g): + return g * np.exp(ctx.saved_data['x']) + + +class _Log(Function): + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x=x) + return np.log(x + 1e-15) + + @staticmethod + def backward(ctx, g): + return g / (ctx.saved_data['x'] + 1e-15) + + +class _Max(Function): + @staticmethod + def forward(ctx, x, axis, keepdims): + ctx.save_for_backward(x=x, axis=axis, keepdims=keepdims) + return np.max(x, axis=axis, keepdims=keepdims) + + @staticmethod + def backward(ctx, g): + x, axis, keepdims = ctx.saved_data['x'], ctx.saved_data['axis'], ctx.saved_data['keepdims'] + mask = (x == np.max(x, axis=axis, keepdims=True)).astype(float) + s = mask.sum(axis=axis, keepdims=True) if axis is not None else mask.sum() + s = np.clip(s, 1e-15, None) + mask = mask / s + if axis is None: + gx = g * mask + else: + gx = np.expand_dims(g, axis=axis) if not keepdims else g + gx = gx * mask + return gx.reshape(x.shape) + + +class _Concatenate(Function): + @staticmethod + def forward(ctx, axis, *tensors): + ctx.save_for_backward(axis=axis, splits=tuple(t.shape[axis] for t in tensors)) + return np.concatenate(tensors, axis=axis) + + @staticmethod + def backward(ctx, g): + axis = ctx.saved_data['axis'] + splits = ctx.saved_data['splits'] + return tuple(np.split(g, np.cumsum(splits[:-1]), axis=axis)) + + +class _Tile(Function): + @staticmethod + def forward(ctx, x, reps): + ctx.save_for_backward(x_shape=x.shape, reps=reps) + return np.tile(x, reps) + + @staticmethod + def backward(ctx, g): + x_shape, reps = ctx.saved_data['x_shape'], ctx.saved_data['reps'] + for ax, r in enumerate(reps): + if r > 1: + g = np.add.reduceat(g, np.arange(0, g.shape[ax], r), axis=ax) + return g.reshape(x_shape) + + +class _Clip(Function): + @staticmethod + def forward(ctx, x, a, b): + ctx.save_for_backward(x=x, a=a, b=b) + return np.clip(x, a, b) + + @staticmethod + def backward(ctx, g): + x, a, b = ctx.saved_data['x'], ctx.saved_data['a'], ctx.saved_data['b'] + return g * ((x >= a) & (x <= b)).astype(float) + + +class _Abs(Function): + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x=x) + return np.abs(x) + + @staticmethod + def backward(ctx, g): + return g * np.sign(ctx.saved_data['x'] + 1e-15) + + +class _Repeat(Function): + @staticmethod + def forward(ctx, x, repeats, axis): + ctx.save_for_backward(x_shape=x.shape, repeats=repeats, axis=axis) + return np.repeat(x, repeats, axis=axis) + + @staticmethod + def backward(ctx, g): + x_shape, repeats, axis = ctx.saved_data['x_shape'], ctx.saved_data['repeats'], ctx.saved_data['axis'] + shape = list(g.shape) + shape.insert(axis + 1, repeats) + shape[axis] = x_shape[axis] + return g.reshape(shape).sum(axis=axis + 1) + + +class _Maximum(Function): + @staticmethod + def forward(ctx, a, b): + ctx.save_for_backward(a=a, b=b) + return np.maximum(a, b) + + @staticmethod + def backward(ctx, g): + a, b = ctx.saved_data['a'], ctx.saved_data['b'] + mask = (a >= b).astype(float) + return g * mask, g * (1.0 - mask) + + def add(a, b): - a, b = _ensure_tensor(a), _ensure_tensor(b) - out = a.data + b.data - result = Tensor(out) - tape = get_active_tape() - if tape and (id(a) in tape._watched or id(b) in tape._watched): - a_shape, b_shape = a.shape, b.shape - def bw(g): - return [broadcast_backward(g, a_shape), - broadcast_backward(g, b_shape)] - tape._record_op([a, b], result, bw, 'add') - return result + return _Add.apply(_ensure_tensor(a), _ensure_tensor(b)) def sub(a, b): - a, b = _ensure_tensor(a), _ensure_tensor(b) - out = a.data - b.data - result = Tensor(out) - tape = get_active_tape() - if tape and (id(a) in tape._watched or id(b) in tape._watched): - a_shape, b_shape = a.shape, b.shape - def bw(g): - return [broadcast_backward(g, a_shape), - broadcast_backward(-g, b_shape)] - tape._record_op([a, b], result, bw, 'sub') - return result + return _Sub.apply(_ensure_tensor(a), _ensure_tensor(b)) def mul(a, b): - a, b = _ensure_tensor(a), _ensure_tensor(b) - out = a.data * b.data - result = Tensor(out) - tape = get_active_tape() - if tape and (id(a) in tape._watched or id(b) in tape._watched): - a_shape, b_shape = a.shape, b.shape - def bw(g): - ga = broadcast_backward(g * b.data, a_shape) if id(a) in tape._watched else None - gb = broadcast_backward(g * a.data, b_shape) if id(b) in tape._watched else None - return [ga, gb] - tape._record_op([a, b], result, bw, 'mul') - return result + return _Mul.apply(_ensure_tensor(a), _ensure_tensor(b)) def div(a, b): - a, b = _ensure_tensor(a), _ensure_tensor(b) - out = a.data / b.data - result = Tensor(out) - tape = get_active_tape() - if tape and (id(a) in tape._watched or id(b) in tape._watched): - a_shape, b_shape = a.shape, b.shape - def bw(g): - ga = broadcast_backward(g / b.data, a_shape) if id(a) in tape._watched else None - gb = broadcast_backward(-g * a.data / (b.data ** 2), b_shape) if id(b) in tape._watched else None - return [ga, gb] - tape._record_op([a, b], result, bw, 'div') - return result + return _Div.apply(_ensure_tensor(a), _ensure_tensor(b)) def neg(x): - x = _ensure_tensor(x) - out = -x.data - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - def bw(g): - return [-g] - tape._record_op([x], result, bw, 'neg') - return result + return _Neg.apply(_ensure_tensor(x)) def _pow(x, power): - x = _ensure_tensor(x) if not isinstance(power, (int, float)): raise ValueError("_pow only supports constant scalar exponent") - out = x.data ** power - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - k = power - def bw(g): - return [(k * (x.data ** (k - 1)) * g)] - tape._record_op([x], result, bw, 'pow') - return result + return _Pow.apply(_ensure_tensor(x), power) def matmul(a, b): - a, b = _ensure_tensor(a), _ensure_tensor(b) - out = a.data @ b.data - result = Tensor(out) - tape = get_active_tape() - if tape and (id(a) in tape._watched or id(b) in tape._watched): - a_shape, b_shape = a.shape, b.shape - def bw(g): - ga = None - if id(a) in tape._watched: - if g.ndim == 2 and b.data.ndim == 2: - ga = g @ b.data.T - else: - ga = g @ np.swapaxes(b.data, -1, -2) - ga = broadcast_backward(ga, a_shape) - gb = None - if id(b) in tape._watched: - if g.ndim == 2 and a.data.ndim == 2: - gb = a.data.T @ g - else: - gb = np.swapaxes(a.data, -1, -2) @ g - gb = broadcast_backward(gb, b_shape) - return [ga, gb] - tape._record_op([a, b], result, bw, 'matmul') - return result + return _Matmul.apply(_ensure_tensor(a), _ensure_tensor(b)) def _sum(x, axis=None, keepdims=False): - x = _ensure_tensor(x) - out = np.sum(x.data, axis=axis, keepdims=keepdims) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - x_shape = x.shape - ax = axis - def bw(g): - if ax is None: - gx = g * np.ones(x_shape, dtype=float) - else: - gx = np.expand_dims(g, axis=ax) if not keepdims else g - gx = gx * np.ones(x_shape, dtype=float) - return [gx.reshape(x_shape)] - tape._record_op([x], result, bw, 'sum') - return result + return _Sum.apply(_ensure_tensor(x), axis=axis, keepdims=keepdims) def _mean(x, axis=None, keepdims=False): - x = _ensure_tensor(x) - out = np.mean(x.data, axis=axis, keepdims=keepdims) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - x_shape = x.shape - ax = axis - if ax is None: - n = x.data.size - else: - ax_t = (ax,) if isinstance(ax, int) else ax - n = int(np.prod([x_shape[d] for d in ax_t])) - def bw(g): - if ax is None: - gx = g * np.ones(x_shape, dtype=float) / n - else: - gx = np.expand_dims(g, axis=ax) if not keepdims else g - gx = gx * np.ones(x_shape, dtype=float) / n - return [gx.reshape(x_shape)] - tape._record_op([x], result, bw, 'mean') - return result + return _Mean.apply(_ensure_tensor(x), axis=axis, keepdims=keepdims) def transpose(x, axes=None): - x = _ensure_tensor(x) - if axes is None: - axes = tuple(range(x.ndim - 1, -1, -1)) - out = np.transpose(x.data, axes) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - inv_axes = tuple(np.argsort(axes)) - def bw(g): - return [np.transpose(g, inv_axes)] - tape._record_op([x], result, bw, 'transpose') - return result + return _Transpose.apply(_ensure_tensor(x), axes) def reshape(x, shape): - x = _ensure_tensor(x) - out = x.data.reshape(shape) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - x_shape = x.shape - def bw(g): - return [g.reshape(x_shape)] - tape._record_op([x], result, bw, 'reshape') - return result + return _Reshape.apply(_ensure_tensor(x), shape) def slice_op(x, idx): - x = _ensure_tensor(x) - out = x.data[idx] - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - def bw(g): - full = np.zeros_like(x.data) - full[idx] += g - return [full] - tape._record_op([x], result, bw, 'slice') - return result + return _Slice.apply(_ensure_tensor(x), idx) def relu(x): - x = _ensure_tensor(x) - out = np.maximum(0, x.data) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - def bw(g): - return [g * (x.data > 0).astype(float)] - tape._record_op([x], result, bw, 'relu') - return result + return _Relu.apply(_ensure_tensor(x)) def sigmoid(x): - x = _ensure_tensor(x) - s = 1.0 / (1.0 + np.exp(-np.clip(x.data, -500, 500))) - out = s.copy() - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - s_cache = s.copy() - def bw(g): - return [g * s_cache * (1.0 - s_cache)] - tape._record_op([x], result, bw, 'sigmoid') - return result + return _Sigmoid.apply(_ensure_tensor(x)) def tanh(x): - x = _ensure_tensor(x) - out = np.tanh(x.data) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - def bw(g): - t = np.tanh(x.data) - return [g * (1.0 - t ** 2)] - tape._record_op([x], result, bw, 'tanh') - return result + return _Tanh.apply(_ensure_tensor(x)) def silu(x): - x = _ensure_tensor(x) - s = 1.0 / (1.0 + np.exp(-np.clip(x.data, -500, 500))) - out = x.data * s - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - s_cache = s.copy() - def bw(g): - ds = s_cache * (1.0 - s_cache) - return [g * (s_cache + x.data * ds)] - tape._record_op([x], result, bw, 'silu') - return result + return _Silu.apply(_ensure_tensor(x)) def softmax(x, axis=-1): - x = _ensure_tensor(x) - x_max = np.max(x.data, axis=axis, keepdims=True) - exps = np.exp(x.data - x_max) - s = exps / np.sum(exps, axis=axis, keepdims=True) - out = s.copy() - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - s_cache = s.copy() - def bw(g): - dot = np.sum(s_cache * g, axis=axis, keepdims=True) - return [s_cache * (g - dot)] - tape._record_op([x], result, bw, 'softmax') - return result + return _Softmax.apply(_ensure_tensor(x), axis=axis) def sqrt(x): - x = _ensure_tensor(x) - out = np.sqrt(x.data) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - def bw(g): - return [g / (2.0 * np.sqrt(x.data) + 1e-15)] - tape._record_op([x], result, bw, 'sqrt') - return result + return _Sqrt.apply(_ensure_tensor(x)) def exp(x): - x = _ensure_tensor(x) - out = np.exp(x.data) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - def bw(g): - return [g * np.exp(x.data)] - tape._record_op([x], result, bw, 'exp') - return result + return _Exp.apply(_ensure_tensor(x)) def log(x): - x = _ensure_tensor(x) - out = np.log(x.data + 1e-15) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - def bw(g): - return [g / (x.data + 1e-15)] - tape._record_op([x], result, bw, 'log') - return result + return _Log.apply(_ensure_tensor(x)) def max_op(x, axis=None, keepdims=False): - x = _ensure_tensor(x) - out = np.max(x.data, axis=axis, keepdims=keepdims) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - x_shape = x.shape - ax = axis - def bw(g): - mask = (x.data == np.max(x.data, axis=ax, keepdims=True)).astype(float) - s = mask.sum(axis=ax, keepdims=True) if ax is not None else mask.sum() - s = np.clip(s, 1e-15, None) - mask = mask / s - if ax is None: - gx = g * mask - else: - gx = np.expand_dims(g, axis=ax) if not keepdims else g - gx = gx * mask - return [gx.reshape(x_shape)] - tape._record_op([x], result, bw, 'max') - return result + return _Max.apply(_ensure_tensor(x), axis=axis, keepdims=keepdims) def concatenate(tensors, axis=0): tensors = [_ensure_tensor(t) for t in tensors] - out = np.concatenate([t.data for t in tensors], axis=axis) - result = Tensor(out) - tape = get_active_tape() - if tape: - splits = [t.shape[axis] for t in tensors] - def bw(g): - grads = np.split(g, np.cumsum(splits[:-1]), axis=axis) - return [g if id(t) in tape._watched else None - for g, t in zip(grads, tensors)] - tape._record_op(tensors, result, bw, 'concatenate') - return result + return _Concatenate.apply(axis, *tensors) def tile(x, reps): - x = _ensure_tensor(x) - out = np.tile(x.data, reps) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - x_shape = x.shape - reps_t = tuple(reps) if isinstance(reps, (tuple, list)) else (reps,) - def bw(g): - for ax, r in enumerate(reps_t): - if r > 1: - g = np.add.reduceat(g, np.arange(0, g.shape[ax], x_shape[ax]), axis=ax) - return [g.reshape(x_shape)] - tape._record_op([x], result, bw, 'tile') - return result + reps = tuple(reps) if isinstance(reps, (tuple, list)) else (reps,) + return _Tile.apply(_ensure_tensor(x), reps) def clip(x, a, b): - x = _ensure_tensor(x) - out = np.clip(x.data, a, b) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - def bw(g): - mask = ((x.data >= a) & (x.data <= b)).astype(float) - return [g * mask] - tape._record_op([x], result, bw, 'clip') - return result + return _Clip.apply(_ensure_tensor(x), a, b) def abs_op(x): - x = _ensure_tensor(x) - out = np.abs(x.data) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - def bw(g): - return [g * np.sign(x.data + 1e-15)] - tape._record_op([x], result, bw, 'abs') - return result - - -def var(x, axis=None, keepdims=False): - mean = _mean(x, axis=axis, keepdims=True) - diff = x - mean - return _mean(diff ** 2, axis=axis, keepdims=keepdims) + return _Abs.apply(_ensure_tensor(x)) def repeat_op(x, repeats, axis): - x = _ensure_tensor(x) - out = np.repeat(x.data, repeats, axis=axis) - result = Tensor(out) - tape = get_active_tape() - if tape and id(x) in tape._watched: - x_shape = x.shape - def bw(g): - shape = list(g.shape) - shape.insert(axis + 1, repeats) - shape[axis] = x_shape[axis] - return [g.reshape(shape).sum(axis=axis + 1)] - tape._record_op([x], result, bw, 'repeat') - return result + return _Repeat.apply(_ensure_tensor(x), repeats, axis) def maximum(a, b): - a, b = _ensure_tensor(a), _ensure_tensor(b) - out = np.maximum(a.data, b.data) - result = Tensor(out) - tape = get_active_tape() - if tape and (id(a) in tape._watched or id(b) in tape._watched): - a_shape, b_shape = a.shape, b.shape - def bw(g): - mask = (a.data >= b.data).astype(float) - ga = broadcast_backward(g * mask, a_shape) if id(a) in tape._watched else None - gb = broadcast_backward(g * (1.0 - mask), b_shape) if id(b) in tape._watched else None - return [ga, gb] - tape._record_op([a, b], result, bw, 'maximum') - return result + return _Maximum.apply(_ensure_tensor(a), _ensure_tensor(b)) + + +def var(x, axis=None, keepdims=False): + mean = _mean(x, axis=axis, keepdims=True) + diff = x - mean + return _mean(diff ** 2, axis=axis, keepdims=keepdims) diff --git a/neutro/autograd/tape.py b/neutro/autograd/tape.py index 1739627..7047300 100644 --- a/neutro/autograd/tape.py +++ b/neutro/autograd/tape.py @@ -20,7 +20,7 @@ def __exit__(self, *args): _tape_stack.pop() def watch(self, tensor): - self._watched.add(id(tensor)) + self._watched.add(tensor) def _record_op(self, inputs, output, backward_fn, name=''): self._ops.append({ @@ -29,29 +29,29 @@ def _record_op(self, inputs, output, backward_fn, name=''): 'backward': backward_fn, 'name': name, }) - self._watched.add(id(output)) + self._watched.add(output) def gradient(self, target, sources): - grad = {id(target): np.ones_like(target.data)} + grad = {target: np.ones_like(target.data)} for op in reversed(self._ops): - out_id = id(op['output']) - if out_id not in grad: + out = op['output'] + if out not in grad: continue - upstream = grad[out_id] + upstream = grad[out] input_grads = op['backward'](upstream) for inp, ig in zip(op['inputs'], input_grads): if ig is not None: - inp_id = id(inp) - if inp_id in grad: - grad[inp_id] += ig + if inp in grad: + grad[inp] += ig else: - grad[inp_id] = ig + grad[inp] = ig + grads = [] for src in sources: - src.grad = grad.get(id(src)) - - return {id(s): grad.get(id(s)) for s in sources} + src.grad = grad.get(src) + grads.append(src.grad) + return grads def reset(self): self._ops.clear() diff --git a/neutro/autograd/tensor.py b/neutro/autograd/tensor.py index 91d1a4e..7a47713 100644 --- a/neutro/autograd/tensor.py +++ b/neutro/autograd/tensor.py @@ -4,7 +4,7 @@ def as_tensor(x): if isinstance(x, Tensor): return x - if isinstance(x, list): + if isinstance(x, (list, tuple)): return [as_tensor(i) for i in x] if isinstance(x, np.ndarray): return Tensor(x) @@ -18,11 +18,16 @@ def __init__(self, data): self.data = np.asarray(data, dtype=float) self.grad = None + __hash__ = object.__hash__ + def __array__(self, dtype=None): if dtype is None: return self.data return self.data.astype(dtype) + def __bool__(self): + return bool(self.data) + def zero_grad(self): self.grad = None @@ -84,11 +89,11 @@ def transpose(self, *axes): axes = axes[0] return transpose(self, axes) - def sum(self, axis=None, keepdims=False, **kwargs): + def sum(self, axis=None, keepdims=False, dtype=None, out=None): from .ops import _sum return _sum(self, axis=axis, keepdims=keepdims) - def mean(self, axis=None, keepdims=False, **kwargs): + def mean(self, axis=None, keepdims=False, dtype=None, out=None): from .ops import _mean return _mean(self, axis=axis, keepdims=keepdims)