Fix transplant, predict, NaN and checkpoint edge cases (3.14.1)

.github/workflows/ci.yml

@@ -19,7 +19,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.10', '3.11', '3.12']
+        python-version: ['3.10', '3.11', '3.12', '3.13']
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5

CHANGELOG.md

@@ -4,6 +4,27 @@ All notable changes to OdyssNet will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
+## [3.14.1] — 2026-10-08
+
+### Fixed
+- **`transplant_weights` left the core diagonal at ~1e-3 after a scale-up.** The `micro_quiet_warm` init it applies before copying the overlap fills the new region's diagonal too, so the first forward pass used self-connections that belong to `memory_feedback`. The diagonal is zeroed once the state is loaded.
+- **`predict` / `evaluate` rejected the integer index tensors `train_batch` accepts** in direct mode, failing on a dtype mismatch. Both now share one input helper.
+- **A non-pulse `forward` kept its input graph on the model**, which held the previous batch's activations until the next call and made `copy.deepcopy(model)` raise. The cache lives only for the call.
+- **A non-finite gradient destroyed the model on CPU.** One NaN target put NaN into ChaosGrad's moments and `W`, and the next good batch could not repair it. The trainer now skips the step with a `RuntimeWarning`, and `GradScaler.update()` still runs so loss-scale backoff is unchanged. Verified on CPU only.
+- **`Neurogenesis.expand` re-measured the traction anchor on the trained weights** (`chaos_core` `rms0` 0.27 → 0.21 and `projections` 0.019 → 0.042 on a 6 → 8 neuron probe), so the cap stopped tracking the initial scale and moved again on every growth. `rms0` is carried over with the other per-group estimates. `verbose=False` no longer prints the "optimizer state transferred" line.
+- **`save_checkpoint` was not atomic.** An interrupted write destroyed the checkpoint it replaced; it now writes beside the target and renames over it. `get_checkpoint_info` counted buffers as parameters (280 against `get_num_params()`'s 132 at 8 neurons with `hebb_type='both'`) and no longer does.
+
+- **`OdyssNetTrainer(model)` moved a GPU model back to the CPU.** The trainer's `device` defaulted to `'cpu'` and the constructor calls `model.to(device)`, so a model built with `device='cuda'` was silently relocated unless the device was repeated. It now defaults to the model's own device; passing `device=` explicitly behaves as before. Read from the code, not run: this machine has no GPU.
+- **Mistyped configuration was accepted silently.** An unknown `weight_init` name fell through to `uniform(-0.1, 0.1)`, a mistyped `vocab_mode` built neither `embed` nor `proj` so the input was ignored, and neuron ids outside `[0, num_neurons)` surfaced as an `IndexError` mid-forward or, for negative ids, wrapped around to the last neuron. Each now raises `ValueError` at construction; `'uniform'` is an explicit strategy.
+
+- **`load_checkpoint(strict=False)` was documented as ignoring size mismatches.** It skips missing and unexpected keys only; a shape mismatch raises, as measured on a 4 → 6 neuron load. The docstring and `docs/LIBRARY.md` now say so and point to `transplant_weights`. The neurogenesis section of `docs/LIBRARY.md` no longer claims incoming weights start at zero: both directions start as `micro_quiet_warm` noise (max 1.2e-3 and 2.5e-3 measured).
+
+### Removed
+- **`pytorch-lightning` and `tqdm` from the dependencies.** Nothing in the package or the examples imports either.
+
+### Changed
+- CI runs on Python 3.13 as well (398 tests pass there). The README quick start falls back to CPU instead of requiring CUDA.
+
 ## [3.14.0] — 2026-10-08
 
 ### Added

CITATION.cff

@@ -5,7 +5,7 @@ authors:
   given-names: "Cahit"
   email: "cksoftwaresystems@gmail.com"
 title: "OdyssNet: The Trainable Dynamic System & Zero-Hidden Architecture"
-version: 3.14.0
+version: 3.14.1
 date-released: 2026-10-08
 url: "https://github.com/theomgdev/OdyssNet"
 abstract: "OdyssNet is a chaotic, fully connected neural network architecture that proves temporal depth (thinking steps) can replace spatial depth (hidden layers). It solves non-linear problems like MNIST with Zero Hidden Layers by utilizing Trainable Chaos."

README.md

@@ -120,8 +120,9 @@ set_seed(42)
 
 # Initialize a Zero-Hidden Network
 # 1 Input, 1 Output. 
-model = OdyssNet(num_neurons=2, input_ids=[0], output_ids=[1], device='cuda')
-trainer = OdyssNetTrainer(model, lr=1e-4, device='cuda')
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model = OdyssNet(num_neurons=2, input_ids=[0], output_ids=[1], device=device)
+trainer = OdyssNetTrainer(model, lr=1e-4, device=device)
 
 # Train
 inputs = torch.randn(100, 1)

README_TR.md

@@ -120,8 +120,9 @@ set_seed(42)
 
 # Sıfır-Gizli Ağ başlat
 # 1 Giriş, 1 Çıkış.
-model = OdyssNet(num_neurons=2, input_ids=[0], output_ids=[1], device='cuda')
-trainer = OdyssNetTrainer(model, lr=1e-4, device='cuda')
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model = OdyssNet(num_neurons=2, input_ids=[0], output_ids=[1], device=device)
+trainer = OdyssNetTrainer(model, lr=1e-4, device=device)
 
 # Eğit
 inputs = torch.randn(100, 1)

docs/LIBRARY.md

@@ -298,6 +298,7 @@ trainer = OdyssNetTrainer(model, optimizer=torch.optim.AdamW(model.parameters(),
 ```
 
 **Parameters:**
+*   `device` (str, optional): Where to train. Default: the model's own device, so a model built with `device='cuda'` stays there.
 *   `lr` (float or None): Learning rate. Default: `None`.
     *   `None`: **ChaosGrad** estimates the step scale online — no manual tuning required. Recommended default. Loss curves vary slightly across runs because the estimate adapts to the observed landscape.
     *   float (e.g. `1e-4`): ChaosGrad runs in **fixed-rate mode** (automatic estimation disabled; AdamW-equivalent updates under the same family policy). Use for byte-for-byte reproducibility studies and benchmarking against fixed baselines.
@@ -336,9 +337,10 @@ Runs a single custom training step. Useful for custom loops (RL, Generative, etc
 *   `full_sequence` (bool): If `True`, calculates loss on the entire sequence output `(Batch, Steps, Out)` instead of just the last step. Essential for Seq2Seq tasks.
 *   `mask` (Tensor, optional): A binary or weighted mask `(Batch, Steps, Out)` to ignore specific steps or outputs during loss calculation. Useful for tasks with "thinking delays" or variable-length sequences.
 *   `output_transform` (Callable, optional): A function to transform the predicted outputs before loss calculation. Useful for reshaping logits (e.g., flatten for CrossEntropy) or applying custom activations.
+*   A step whose gradient norm is not finite is skipped with a `RuntimeWarning`: weights and optimizer state are left as they were.
 
 #### `trainer.predict(input_features, thinking_steps, full_sequence=False)`
-Runs inference in evaluation mode.
+Runs inference in evaluation mode. Takes the same inputs as `train_batch`, integer index tensors included.
 *   `full_sequence` (bool): If `True`, returns outputs for all time steps `(Batch, Steps, Out)`.
 
 #### `trainer.regenerate_synapses(threshold=0.01)`
@@ -897,10 +899,10 @@ The `odyssstore` module provides checkpoint management utilities, including a un
 ### Functions
 
 #### `save_checkpoint(model, optimizer, epoch, loss, path, extra_data=None, trainer_state=None)`
-Saves a training checkpoint to disk. Pass `trainer_state=trainer.state_dict()` to also persist the trainer's runtime state (step counter, scaler, persistent gradients).
+Saves a training checkpoint to disk. Pass `trainer_state=trainer.state_dict()` to also persist the trainer's runtime state (step counter, scaler, persistent gradients). The file is written beside `path` and renamed over it, so an interrupted save leaves the previous checkpoint intact.
 
 #### `load_checkpoint(model, optimizer, path, device='cpu', strict=True, lr=None, trainer=None)`
-Loads a checkpoint. Set `strict=False` to ignore size mismatches (will partially load what fits). Pass `lr` to overwrite the saved learning rate after loading. Pass `trainer` (an `OdyssNetTrainer` instance) to restore runtime trainer state (step counter, scaler, persistent gradients).
+Loads a checkpoint. Set `strict=False` to skip missing or unexpected keys; a size mismatch raises either way, so use `transplant_weights` to load across sizes. Pass `lr` to overwrite the saved learning rate after loading. Pass `trainer` (an `OdyssNetTrainer` instance) to restore runtime trainer state (step counter, scaler, persistent gradients).
 
 #### `transplant_weights(model, checkpoint_path, device='cpu', verbose=True)`
 🧬 **Weight Transplantation**: Transfers learned weights from a checkpoint to a model, **even if the number of neurons is different**.
@@ -934,17 +936,14 @@ OdyssNet supports dynamic growth, allowing you to add neurons to a live network
 Dynamically adds `amount` empty neurons to the model.
 *   **Continuity**: Optimizers are migrated, so momentum and history are preserved.
 *   **State**: The training state is preserved.
-*   **Initialization**: 
-    *   **Incoming Weights**: 0 (Maintains forward pass stability, new neuron starts inactive).
-    *   **Outgoing Weights**: Small random noise (Enables backpropagation / gradient flow).
 
 ```python
 # Add 1 neuron if loss stagnates
 if loss > prev_loss:
     trainer.expand(amount=1)
 ```
 
-> **Initialization:** New connections are initialized with `micro_quiet_warm` (Normal(0, 1e-3)) noise so they remain dormant relative to trained weights and do not destabilize the existing dynamics. Optimizer momentum is migrated from the old parameters to the expanded ones.
+> **Initialization:** Connections to and from a new neuron are initialized with `micro_quiet_warm` (Normal(0, 1e-3)) noise, in both directions, so they remain dormant relative to trained weights and do not destabilize the existing dynamics. Optimizer momentum is migrated from the old parameters to the expanded ones.
 
 ---
 

odyssnet/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "3.14.0"
+__version__ = "3.14.1"
 
 from .core.network import OdyssNet
 from .core.attention import TemporalAttention

odyssnet/core/network.py

@@ -24,6 +24,12 @@ def __init__(self, num_neurons, input_ids, output_ids, pulse_mode=True, dropout_
             else:
                 num_neurons = 0
 
+        for name, ids in (("input_ids", input_ids), ("output_ids", output_ids)):
+            if any(not 0 <= i < num_neurons for i in ids):
+                raise ValueError(f"{name} must lie in [0, {num_neurons}), got {list(ids)}")
+        if vocab_mode not in ('hybrid', 'discrete', 'continuous'):
+            raise ValueError(f"vocab_mode must be 'hybrid', 'discrete' or 'continuous', got {vocab_mode!r}")
+
         self.num_neurons = num_neurons
         self.debug = debug
         if debug:
@@ -79,7 +85,6 @@ def __init__(self, num_neurons, input_ids, output_ids, pulse_mode=True, dropout_
 
         self.pulse_mode = pulse_mode
         self.gradient_checkpointing = gradient_checkpointing
-        self._cached_scaled_input = None
 
         # Parse configurable component settings
         weight_init = self._normalize_weight_init(weight_init)
@@ -394,8 +399,10 @@ def _apply_init(self, tensor, strategy):
                         frob = tensor.norm()
                         if frob > 1e-8:
                             tensor.div_(frob / (tensor.numel() ** 0.5))
-            else:
+            elif strategy == 'uniform':
                 nn.init.uniform_(tensor, -0.1, 0.1)
+            else:
+                raise ValueError(f"Unknown weight init strategy: {strategy!r}")
 
     def regenerate_weak_weights(self, threshold=0.01, percentage=None):
         with torch.no_grad():
@@ -647,6 +654,11 @@ def _plastic_step(h_t_in, local_mem, t_idx, x_input_info,
         if len(output_pos) > 0:
             output_scale_vec[output_pos] = self._get_output_scale(h_t.dtype)
 
+        # Held across steps for continuous (non-pulse) injection; a local, so
+        # the previous call's graph is not kept alive and the model stays
+        # deep-copyable.
+        cached_input = None
+
         for t in range(steps):
             # Prepare input for this step
             x_step_info = None
@@ -702,11 +714,11 @@ def _plastic_step(h_t_in, local_mem, t_idx, x_input_info,
 
                         # Cache input for continuous (non-pulse) persistence across steps
                         if not self.pulse_mode:
-                            self._cached_scaled_input = x_step_info
+                            cached_input = x_step_info
 
                     # Re-use cached input for non-active steps in continuous mode
                     if not self.pulse_mode and x_step_info is None:
-                        x_step_info = getattr(self, '_cached_scaled_input', None)
+                        x_step_info = cached_input
 
                 # --- LEGACY DIRECT MODE ---
                 else:
@@ -759,8 +771,8 @@ def _plastic_step(h_t_in, local_mem, t_idx, x_input_info,
                     else:
                         # Continuous mode: cache on first step, reuse for all subsequent steps
                         if t == 0:
-                            self._cached_scaled_input = x_input * input_scale_vec
-                        x_step_info = self._cached_scaled_input
+                            cached_input = x_input * input_scale_vec
+                        x_step_info = cached_input
 
 
             if self.attn is not None and (self.attn_read == 'step' or t % ratio == 0):

odyssnet/training/trainer.py

@@ -4,14 +4,15 @@
 import time
 import math
 import numbers
+import warnings
 from typing import Callable
 from ..utils.data import prepare_input, to_tensor
 
 from ..utils.neurogenesis import Neurogenesis
 from .chaos_optimizer import ChaosGrad
 
 class OdyssNetTrainer:
-    def __init__(self, model, optimizer=None, loss_fn=None, lr=None, device='cpu',
+    def __init__(self, model, optimizer=None, loss_fn=None, lr=None, device=None,
                  gradient_persistence=0.0, synaptic_noise=0.0,
                  anomaly_hook=None):
         """
@@ -28,14 +29,14 @@ def __init__(self, model, optimizer=None, loss_fn=None, lr=None, device='cpu',
                 - float (e.g. 1e-4): ChaosGrad runs with a fixed AdamW-style
                   learning rate (automatic estimation disabled). Use for
                   exact reproducibility studies.
-            device (str): Device to run training on.
+            device (str): Device to run training on. Default: the model's own device.
             gradient_persistence (float): How much gradient to keep from previous step (0.0-0.9).
             synaptic_noise (float): Scale of noise added to weights during training. Default 0.0.
             anomaly_hook (callable, optional): Called as hook(event_type, loss_value) on
                 anomalies ('spike', 'plateau', 'increase').
         """
         self.model = model
-        self.device = device
+        self.device = device if device is not None else model.device
         self.model.to(self.device)
         self.gradient_persistence = gradient_persistence
         self.synaptic_noise = synaptic_noise
@@ -105,6 +106,21 @@ def _get_autocast_ctx(self):
             device_type='cuda' if is_cuda else 'cpu', enabled=is_cuda,
         )
 
+    def _prepare_inputs(self, input_features):
+        """Map user input to what `forward` takes, and return its batch size.
+
+        Vocab models take token ids or raw vectors as they are. Without a
+        vocab, integer tensors are neuron indices to inject, and anything else
+        is a feature matrix mapped onto `input_ids`.
+        """
+        if getattr(self.model, 'vocab_size', None) is not None:
+            x_input = to_tensor(input_features, self.device)
+        elif isinstance(input_features, torch.Tensor) and input_features.dtype in (torch.int32, torch.int64):
+            x_input = input_features.to(self.device)
+        else:
+            return prepare_input(input_features, self.model.input_ids, self.model.num_neurons, self.device)
+        return x_input, x_input.shape[0]
+
     def _extract_outputs(self, all_states, final_state, full_sequence):
         """Extract the prediction tensor from forward-pass outputs.
 
@@ -205,17 +221,7 @@ def train_batch(self, input_features, target_values, thinking_steps, gradient_ac
                         if 'W' in name and param.dim() == 2 and param.shape[0] == param.shape[1]:
                             param.fill_diagonal_(0.0)
 
-        # Prepare Data
-        # If model has vocab_size, we assume input is Token IDs or Raw Vects for Projection.
-        # We bypass 'prepare_input' which attempts to map features to specific neurons manually.
-        if hasattr(self.model, 'vocab_size') and self.model.vocab_size is not None:
-            x_input = to_tensor(input_features, self.device)
-            batch_size = x_input.shape[0]
-        elif isinstance(input_features, torch.Tensor) and input_features.dtype in [torch.long, torch.int, torch.int32, torch.int64]:
-            x_input = input_features.to(self.device)
-            batch_size = x_input.shape[0]
-        else:
-            x_input, batch_size = prepare_input(input_features, self.model.input_ids, self.model.num_neurons, self.device)
+        x_input, batch_size = self._prepare_inputs(input_features)
 
         target_values = to_tensor(target_values, self.device)
         if mask is not None:
@@ -260,8 +266,13 @@ def train_batch(self, input_features, target_values, thinking_steps, gradient_ac
             if self.gradient_persistence > 0.0:
                 self._inject_persistent_grads()
 
-            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
-            self.scaler.step(self.optimizer)
+            grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+            # A non-finite gradient would poison the optimizer's moments and
+            # the weights for good; GradScaler only guards this under CUDA AMP.
+            if torch.isfinite(grad_norm):
+                self.scaler.step(self.optimizer)
+            else:
+                warnings.warn("Non-finite gradient; optimizer step skipped.", RuntimeWarning, stacklevel=2)
             self.scaler.update()
 
             if self.gradient_persistence > 0.0:
@@ -353,12 +364,7 @@ def predict(self, input_features, thinking_steps, full_sequence=False):
         """
         self.model.eval()
         with torch.no_grad():
-            if hasattr(self.model, 'vocab_size') and self.model.vocab_size is not None:
-                x_input = to_tensor(input_features, self.device)
-                batch_size = x_input.shape[0]
-            else:
-                x_input, batch_size = prepare_input(input_features, self.model.input_ids, self.model.num_neurons, self.device)
-
+            x_input, batch_size = self._prepare_inputs(input_features)
             self.model.reset_state(batch_size)
             all_states, final_state = self.model(x_input, steps=thinking_steps, return_sequence=full_sequence)
             return self._extract_outputs(all_states, final_state, full_sequence)