[StackRnn] - add unrolled (own-weights) mode; each position gets its own RBM

Previously a single shared-weight Entity processed every time step.  Now:
- Shared mode  (1 layer via make_layer):  original behaviour unchanged
- Unrolled mode (N layers via make_unrolled): layers[t] owns W_t, b_v_t, b_h_t

New API:
  make_unrolled(T, sensory_size, h_size, ...)  → list[Layer]
  next_entity()   → Entity for the upcoming step() call
  current_entity() → Entity from the most recent step() call
  is_shared        → bool

moby_rnn.ipynb: switch Build model cell to make_unrolled(T=100); update
  predict_next() to use rnn.next_entity() for position-correct Gibbs sampling
README_moby_rnn.md: redraw temporal-unrolling ASCII art showing per-position
  weights W_t; update parameter count and checkpoint file listing

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

README_moby_rnn.md

@@ -38,26 +38,31 @@ of a **context** vector (the previous hidden state) and the current sensory inpu
    (next time step)     (predict char)
 ```
 
-### Temporal unrolling
+### Temporal unrolling (unrolled / own-weights mode)
 
 ```
       x_{t-1}              x_t               x_{t+1}
          │                   │                   │
-  ┌──────┴──────┐  ┌──────┴──────┐  ┌──────┴──────┐
+  ┌──────┴──────┐    ┌───────┴──────┐    ┌───────┴──────┐
   │  ctx  │x_{t-1}  │  ctx  │  x_t  │   │  ctx  │x_{t+1}│
-  │       ╠═══════╣ │       ╠═══════╣ │       ╠═══════╣
+  │       ╠═══════╣  │       ╠════════╣   │       ╠════════╣
-  │    W (shared)   │    W (shared)   │    W (shared)   │
+  │    W_{t-1}    │  │      W_t      │   │     W_{t+1}   │
-  │       ╠═══════╣ │       ╠═══════╣ │       ╠═══════╣
+  │       ╠═══════╣  │       ╠════════╣   │       ╠════════╣
-  │   hidden_t-1  │ │    hidden_t   │ │   hidden_t+1  │
+  │  hidden_{t-1} │  │   hidden_t    │   │  hidden_{t+1} │
-  └───────┬───────┘ └───────┬───────┘ └───────┬───────┘
+  └───────┬───────┘  └───────┬────────┘  └───────┬────────┘
-          │   context_t-1   │   context_t      │   context_t+1
+          │  context_{t-1}   │  context_t         │  context_{t+1}
-          └────────────────►┘ ────────────────►┘ ──────► ...
+          └─────────────────►┘ ──────────────────►┘ ──────► ...
 ```
 
-Weights **W**, **b_v**, **b_h** are shared across all time steps — the same RBM
+Each time step has its **own weight matrix** W_t, b_v_t, b_h_t — the model is
-processes every character.  This is the RTRBM concatenation variant: instead of
+**unrolled**: N positions in a sequence → N separate RBMs.  This lets each
-modulating the hidden biases (Sutskever & Hinton, 2007), the previous hidden
+position specialise for the statistical patterns that occur at that offset in a
-state is directly concatenated to the visible layer.
+sequence.  The sequence length T must be fixed and equal to N at training time;
+during generation, position indices wrap modulo N.
+
+For comparison, the shared-weights variant (1 entity, `make_layer`) uses a
+single W across all time steps — the RTRBM concatenation variant of Sutskever &
+Hinton (2007).
 
 ---
 
@@ -78,7 +83,8 @@ All hyperparameters live in the **Build model** cell of `moby_rnn.ipynb`:
 | `SEED` | `"Call me Ishmael."` | Seed text for free generation |
 | `SEQ_IDX` | 42 | Sequence index for the hidden-state trace plot |
 
-Model size with defaults: **(512 + 85) × 512 = 306,176 parameters**.
+Model size with defaults: **(CONTEXT_SIZE + 85) × CONTEXT_SIZE × T parameters** —
+e.g. T=100, CONTEXT_SIZE=256 → 100 × 27,136 = **2,713,600 parameters** total.
 
 ---
 
@@ -107,7 +113,10 @@ continue training from the last saved state — `model.state_load()` in cell 4
 picks it up automatically.
 
 ```
-results/moby_rnn-0-state.npz   ← single-layer checkpoint
+results/moby_rnn-0-state.npz      ← RBM at position 0
+results/moby_rnn-1-state.npz      ← RBM at position 1
+...
+results/moby_rnn-99-state.npz     ← RBM at position T-1
 ```
 
 ---

moby_rnn.ipynb

@@ -158,48 +158,9 @@
      "start_time": "2026-05-31T09:38:49.498420274Z"
     }
    },
-   "source": [
+   "source": "# ── Build model ────────────────────────────────────────────────────────────\n# Unrolled mode: T separate RBMs, one per position in the sequence.\n# Each RBM_t has its own W_t, b_v_t, b_h_t.\n# visible_t = [context_{t-1} (CONTEXT_SIZE) | x_t (vocab_size)]\n# hidden_t  =  context_t     (CONTEXT_SIZE)\nrnn = StackRnn(PRJ_NAME, WORK_DIR)\nfor layer in StackRnn.make_unrolled(\n    T, vocab_size, CONTEXT_SIZE,\n    entity_params=EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False),\n    training_params=TrainingParams(\n        learning_rate=0.005,\n        momentum=0.9,\n        num_epochs=200,\n        do_rao_blackwell=True,\n        l2_lambda=0.0001,\n    ),\n):\n    rnn.append(layer)\n\nrnn.state_init(0.01)\nrnn.state_load()   # resumes from checkpoint if one exists\n\ne = rnn.from_index(0).entity\nprint(f\"Mode            : unrolled ({rnn.num_layers()} layers, own weights per position)\")\nprint(f\"Visible         : {rnn.h_size()} (context) + {rnn.sensory_size()} (vocab) = {e.shape[0]}\")\nprint(f\"Hidden          : {rnn.h_size()}\")\nprint(f\"Params/layer    : {e.shape[0] * e.shape[1]:,}\")\nprint(f\"Total params    : {e.shape[0] * e.shape[1] * rnn.num_layers():,}\")",
-    "# ── Build model ────────────────────────────────────────────────────────────\n",
+   "outputs": [],
-    "# visible layer = [context (CONTEXT_SIZE) | x_t (vocab_size)]\n",
+   "execution_count": null
-    "# hidden layer  = h_t (CONTEXT_SIZE)\n",
-    "rnn = StackRnn(PRJ_NAME, WORK_DIR)\n",
-    "rnn.append(StackRnn.make_layer(\n",
-    "    \"layer0\",\n",
-    "    sensory_size=vocab_size,\n",
-    "    h_size=CONTEXT_SIZE,\n",
-    "    entity_params=EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False),\n",
-    "    training_params=TrainingParams(\n",
-    "        learning_rate=0.005,\n",
-    "        momentum=0.9,\n",
-    "        num_epochs=200,\n",
-    "        do_rao_blackwell=True,\n",
-    "        l2_lambda=0.0001,\n",
-    "    ),\n",
-    "))\n",
-    "\n",
-    "rnn.state_init(0.01)\n",
-    "rnn.state_load()   # resumes from checkpoint if one exists\n",
-    "\n",
-    "e = rnn.from_index(0).entity\n",
-    "print(f\"Entity          : {e.name}\")\n",
-    "print(f\"Visible         : {rnn.h_size()} (context) + {rnn.sensory_size()} (vocab) = {e.shape[0]}\")\n",
-    "print(f\"Hidden          : {rnn.h_size()}\")\n",
-    "print(f\"Parameters      : {e.shape[0] * e.shape[1]:,}\")"
-   ],
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "results/moby_rnn-0-state.npz loaded successfully!\n",
-      "Entity          : Entity-213x128\n",
-      "Visible         : 128 (context) + 85 (vocab) = 213\n",
-      "Hidden          : 128\n",
-      "Parameters      : 27,264\n"
-     ]
-    }
-   ],
-   "execution_count": 58
   },
   {
    "cell_type": "code",
@@ -349,17 +310,9 @@
      "start_time": "2026-05-31T09:39:43.947416428Z"
     }
    },
-   "source": "# ── Generation helpers ─────────────────────────────────────────────────────\n\ndef predict_next(rnn, context, n_gibbs=10, temperature=1.0):\n    \"\"\"Return a character probability distribution conditioned on context.\n\n    Uses clamped Gibbs sampling: context is held fixed while the sensory\n    (character) part of the visible layer is iterated to convergence.\n    \"\"\"\n    entity = rnn.from_index(0).entity\n    h_sz   = rnn.h_size()\n    s_sz   = rnn.sensory_size()\n\n    x_init  = (np.random.rand(1, s_sz) > 0.5).astype(float)\n    visible = np.concatenate([context, x_init], axis=1)\n\n    for _ in range(n_gibbs):\n        h       = entity.forward(visible)\n        visible = entity.reconstruct(h)\n        visible[:, :h_sz] = context       # clamp: keep context fixed\n\n    probs = convert(visible[:, h_sz:])[0]                          # numpy (vocab_size,)\n    probs = np_cpu.power(np_cpu.clip(probs, 1e-10, 1.0), 1.0 / temperature)\n    probs /= probs.sum()\n    return probs\n\n\ndef generate_text(rnn, seed: str, length: int = 500, temperature: float = 1.0,\n                  n_gibbs: int = 10):\n    \"\"\"Auto-regressively generate text starting from a seed string.\"\"\"\n    rnn.reset(batch_size=1)\n    h = np.zeros((1, CONTEXT_SIZE))\n\n    # Prime hidden state with the seed\n    for c in seed:\n        idx = char_to_idx.get(c, 0)\n        x   = np.zeros((1, vocab_size))\n        x[0, idx] = 1.0\n        h   = rnn.step(x)\n\n    generated = seed\n    for _ in range(length):\n        probs  = predict_next(rnn, h.copy(), n_gibbs=n_gibbs, temperature=temperature)\n        idx    = int(np_cpu.random.choice(vocab_size, p=probs))\n        c      = idx_to_char[idx]\n        generated += c\n\n        x         = np.zeros((1, vocab_size))\n        x[0, idx] = 1.0\n        h         = rnn.step(x)\n\n    return generated\n\nprint(\"Helpers defined.\")",
+   "source": "# ── Generation helpers ─────────────────────────────────────────────────────\n\ndef predict_next(rnn, context, n_gibbs=10, temperature=1.0):\n    \"\"\"Return a character probability distribution conditioned on context.\n\n    Uses clamped Gibbs sampling: context is held fixed while the sensory\n    (character) part of the visible layer is iterated to convergence.\n    Uses rnn.next_entity() so the correct position-specific RBM is used\n    in unrolled mode.\n    \"\"\"\n    entity = rnn.next_entity()\n    h_sz   = rnn.h_size()\n    s_sz   = rnn.sensory_size()\n\n    x_init  = (np.random.rand(1, s_sz) > 0.5).astype(float)\n    visible = np.concatenate([context, x_init], axis=1)\n\n    for _ in range(n_gibbs):\n        h       = entity.forward(visible)\n        visible = entity.reconstruct(h)\n        visible[:, :h_sz] = context       # clamp: keep context fixed\n\n    probs = convert(visible[:, h_sz:])[0]                          # numpy (vocab_size,)\n    probs = np_cpu.power(np_cpu.clip(probs, 1e-10, 1.0), 1.0 / temperature)\n    probs /= probs.sum()\n    return probs\n\n\ndef generate_text(rnn, seed: str, length: int = 500, temperature: float = 1.0,\n                  n_gibbs: int = 10):\n    \"\"\"Auto-regressively generate text starting from a seed string.\"\"\"\n    rnn.reset(batch_size=1)\n    h = np.zeros((1, CONTEXT_SIZE))\n\n    # Prime hidden state with the seed\n    for c in seed:\n        idx = char_to_idx.get(c, 0)\n        x   = np.zeros((1, vocab_size))\n        x[0, idx] = 1.0\n        h   = rnn.step(x)\n\n    generated = seed\n    for _ in range(length):\n        probs  = predict_next(rnn, h.copy(), n_gibbs=n_gibbs, temperature=temperature)\n        idx    = int(np_cpu.random.choice(vocab_size, p=probs))\n        c      = idx_to_char[idx]\n        generated += c\n\n        x         = np.zeros((1, vocab_size))\n        x[0, idx] = 1.0\n        h         = rnn.step(x)\n\n    return generated\n\nprint(\"Helpers defined.\")",
-   "outputs": [
+   "outputs": [],
-    {
+   "execution_count": null
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Helpers defined.\n"
-     ]
-    }
-   ],
-   "execution_count": 60
   },
   {
    "cell_type": "code",

src/rbm/stack_rnn.py

@@ -16,33 +16,65 @@ _CD_FUNC = {
 
 
 class StackRnn(Stack):
-    """Recurrent RBM stack (RTRBM-style, concatenation variant).
+    """Recurrent RBM stack.
 
-    At each time step t and layer i:
+    Two modes selected by the number of appended layers:
-        visible_i[t] = concat(h_i[t-1], input_i[t])
-        h_i[t]       = entity_i.forward(visible_i[t])
 
-    where input_0[t] = x_t (sensory input) and input_{i+1}[t] = h_i[t].
+    Shared weights  (1 layer):
-    Weights are shared across time — the same Entity processes every time step.
+        The same Entity processes every time step.
+        Sequences may have any length.
+
+    Unrolled / own weights  (N layers):
+        layers[t % N] processes time step t — each position has its own W.
+        Training requires sequences of exactly length N.
     """
 
     def __init__(self, name: str, work_dir: str = '.'):
         Stack.__init__(self, StackType.Rnn, name, work_dir)
-        self._h: list[Mat] | None = None
+        self._h: Mat | None = None  # context / hidden state: (batch, h_size)
+        self._t: int = 0             # current time-step counter
 
-    # ── Convenience factory ───────────────────────────────────────────────────
+    # ── Convenience factories ─────────────────────────────────────────────
 
     @staticmethod
     def make_layer(name: str, sensory_size: int, h_size: int,
                    entity_params: EntityParams, training_params: TrainingParams) -> Layer:
-        """Create a Layer suitable for StackRnn.
+        """Single layer for shared-weights mode."""
-
-        The entity's visible size = sensory_size + h_size.
-        sensory_size is the raw input width; h_size is the recurrent state size.
-        """
         return Layer(name, (1, sensory_size, h_size, h_size), entity_params, training_params)
 
-    # ── Derived sizes ─────────────────────────────────────────────────────────
+    @staticmethod
+    def make_unrolled(time_steps: int, sensory_size: int, h_size: int,
+                      entity_params: EntityParams, training_params: TrainingParams) -> list[Layer]:
+        """N layers for unrolled (own-weights) mode — one per time step.
+
+        Usage::
+
+            for layer in StackRnn.make_unrolled(T, ...):
+                rnn.append(layer)
+        """
+        return [
+            Layer(f"t{t}", (1, sensory_size, h_size, h_size), entity_params, training_params)
+            for t in range(time_steps)
+        ]
+
+    # ── Mode ─────────────────────────────────────────────────────────────
+
+    @property
+    def is_shared(self) -> bool:
+        return self.num_layers() == 1
+
+    def _entity_at(self, t: int) -> Entity:
+        return self.from_index(t % self.num_layers()).entity
+
+    def next_entity(self) -> Entity:
+        """Entity that will be used by the next step() call."""
+        return self._entity_at(self._t)
+
+    def current_entity(self) -> Entity:
+        """Entity used by the most recent step() call."""
+        return self._entity_at(max(0, self._t - 1))
+
+    # ── Derived sizes ─────────────────────────────────────────────────────
 
     def h_size(self, layer_idx: int = 0) -> int:
         return self.from_index(layer_idx).entity.shape[1]
@@ -51,55 +83,55 @@ class StackRnn(Stack):
         e = self.from_index(layer_idx).entity
         return e.shape[0] - e.shape[1]
 
-    # ── Hidden-state management ───────────────────────────────────────────────
+    # ── State management ──────────────────────────────────────────────────
 
     def reset(self, batch_size: int = 1):
-        """Zero all hidden states (call before processing a new sequence)."""
+        """Zero the context vector and reset the time-step counter."""
-        self._h = [np.zeros((batch_size, self.h_size(i))) for i in range(self.num_layers())]
+        self._h = np.zeros((batch_size, self.h_size()))
+        self._t = 0
 
-    # ── Inference ─────────────────────────────────────────────────────────────
+    # ── Inference ─────────────────────────────────────────────────────────
 
     def step(self, x: Mat) -> Mat:
-        """One time step forward through all layers.
+        """One time step forward.
+
+        In shared mode  uses the single entity.
+        In unrolled mode uses layers[_t % N] and advances _t.
 
         x: (batch_size, sensory_size) or (sensory_size,)
-        Returns the top-layer hidden state h.
+        Returns the new context vector h_t.
         """
         if x.ndim == 1:
             x = x[None, :]
         batch_size = x.shape[0]
-        if self._h is None or self._h[0].shape[0] != batch_size:
+        if self._h is None or self._h.shape[0] != batch_size:
             self.reset(batch_size)
 
-        x_in = _to_gpu(x)
+        entity  = self._entity_at(self._t)
-        for i, layer in enumerate(self.layers):
+        visible = np.concatenate([self._h, _to_gpu(x)], axis=1)
-            visible = np.concatenate([self._h[i], x_in], axis=1)
+        self._h = entity.forward(visible)
-            h_new = layer.entity.forward(visible)
+        self._t += 1
-            self._h[i] = h_new
+        return self._h
-            x_in = h_new
 
-        return self._h[-1]
+    def reconstruct(self, h: Mat) -> Mat:
-
-    def reconstruct(self, h: Mat, layer_idx: int = -1) -> Mat:
         """Decode h → visible, returning only the sensory portion.
 
-        The visible layer is [context | sensory]; this method strips context,
+        Uses the entity from the most recent step() call.
-        returning only the sensory reconstruction.
-        h: (batch_size, h_size)
         """
-        if layer_idx < 0:
+        entity  = self.current_entity()
-            layer_idx = self.num_layers() + layer_idx
-        entity = self.from_index(layer_idx).entity
         visible = entity.reconstruct(h)
         return visible[:, entity.shape[1]:]
 
-    # ── Training ──────────────────────────────────────────────────────────────
+    # ── Training ──────────────────────────────────────────────────────────
 
     def train(self, sequences: Mat, status: Status = None):
-        """Greedy layer-wise CD training over sequences.
+        """Train on sequences.
 
         sequences: (T, sensory_size)           — single sequence
                    (num_seq, T, sensory_size)  — batch of sequences
+
+        Shared mode   (1 layer): T may be any value.
+        Unrolled mode (N layers): T must equal N.
         """
         if status is None:
             status = Status()
@@ -107,18 +139,25 @@ class StackRnn(Stack):
         seqs = sequences if sequences.ndim == 3 else sequences[None, :]
         num_seq, T, _ = seqs.shape
 
-        for layer_idx, layer in enumerate(self.layers):
+        if self.is_shared:
-            entity = layer.entity
+            entity = self.from_index(0).entity
-            print(f"Train layer {layer_idx} ({entity.name}) for {entity.training_params.num_epochs} epochs")
+            print(f"Train shared ({entity.name}) "
-
+                  f"for {entity.training_params.num_epochs} epochs")
             if entity.enable_training and entity.training_params is not None:
-                self._train_layer(entity, seqs, num_seq, T, status)
+                self._train_shared(entity, seqs, num_seq, T, status)
+        else:
+            assert T == self.num_layers(), (
+                f"Unrolled mode: sequence length T={T} "
+                f"must equal num_layers={self.num_layers()}"
+            )
+            params = self.from_index(0).entity.training_params
+            print(f"Train unrolled ({self.num_layers()} layers) "
+                  f"for {params.num_epochs} epochs")
+            self._train_unrolled(seqs, num_seq, T, status)
 
-            if layer_idx < self.num_layers() - 1:
+    def _train_shared(self, entity: Entity, seqs: Mat,
-                seqs = self._pass_through(entity, seqs, num_seq, T)
-
-    def _train_layer(self, entity: Entity, seqs: Mat,
                       num_seq: int, T: int, status: Status):
+        """One entity, reused at every time step."""
         cd_func = _CD_FUNC[entity.type]
         params  = entity.training_params
         h_sz    = entity.shape[1]
@@ -157,7 +196,6 @@ class StackRnn(Stack):
                     keep_running = False
                     break
 
-        # Final report using one clean forward pass
         h = np.zeros((num_seq, h_sz))
         err_total = 0.0
         for t in range(T):
@@ -170,15 +208,57 @@ class StackRnn(Stack):
             "err_rms_total": {"value": err_total / T, "unit": ""},
         })
 
-    def _pass_through(self, entity: Entity, seqs: Mat,
+    def _train_unrolled(self, seqs: Mat, num_seq: int, T: int, status: Status):
-                      num_seq: int, T: int) -> _np_cpu.ndarray:
+        """N entities, one per time step — each has its own W, b_v, b_h."""
-        """Run sequences through entity; return hidden outputs on CPU."""
+        params       = self.from_index(0).entity.training_params
-        h_sz = entity.shape[1]
+        h_sz         = self.from_index(0).entity.shape[1]
-        outputs = _np_cpu.zeros((num_seq, T, h_sz), dtype=_np_cpu.float64)
+        d_progress   = 100.0 / params.num_epochs
+        progress     = 0.0
+        keep_running = True
+
+        for layer in self.layers:
+            layer.entity.grad_zero()
+        status.on_change(self.from_index(0).entity)
+
+        for epoch in range(params.num_epochs):
+            if not keep_running:
+                break
+
             h = np.zeros((num_seq, h_sz))
-        for t in range(T):
+            err_total = 0.0
+
+            for t, layer in enumerate(self.layers):
+                entity  = layer.entity
+                cd_func = _CD_FUNC[entity.type]
+
+                x_t     = _to_gpu(seqs[:, t, :])
+                visible = np.concatenate([h, x_t], axis=1)
+
+                dwhv, dbv, dbh = cd_func(entity, visible)
+                grad = entity.grad_compute(dbv, dbh, dwhv)
+                entity.state_adjust(grad, 1.0 / num_seq)
+
+                h = entity.forward(visible)
+                err_total += rms_error_accu(visible - entity.reconstruct(h))
+
+            progress += d_progress
+            if status.want_report(round(progress)):
+                if not status.on_change(self.from_index(0).entity, {
+                    "progress": {"value": round(progress), "unit": "%"},
+                    "err_rms":  {"value": err_total / T,   "unit": ""},
+                }):
+                    keep_running = False
+                    break
+
+        h = np.zeros((num_seq, h_sz))
+        err_total = 0.0
+        for t, layer in enumerate(self.layers):
+            entity  = layer.entity
             x_t     = _to_gpu(seqs[:, t, :])
             visible = np.concatenate([h, x_t], axis=1)
             h = entity.forward(visible)
-            outputs[:, t, :] = convert(h)
+            err_total += rms_error_accu(visible - entity.reconstruct(h))
-        return outputs
+        status.on_change(self.from_index(0).entity, {
+            "progress":      {"value": 100,           "unit": "%"},
+            "err_rms_total": {"value": err_total / T, "unit": ""},
+        })

src/tests/test_rnn.py

@@ -1,8 +1,4 @@
-"""Test for StackRnn: recurrent RBM with concatenated [h_{t-1} | x_t] visible layer.
+"""Tests for StackRnn — shared-weights and unrolled (own-weights) modes."""
-
-Uses a repeating binary pattern as a minimal synthetic sequence so the model
-has something learnable to compress and predict.
-"""
 import numpy as _np_cpu
 from rbm.stack_rnn import StackRnn
 from rbm.matrix import np, convert
@@ -15,89 +11,80 @@ T = 16           # sequence length
 NUM_SEQ      = 10    # sequences in the batch
 WORK_DIR     = "../../results"
 
+_PARAMS = TrainingParams(learning_rate=0.05, momentum=0.5, num_epochs=5,
+                         do_rao_blackwell=True)
+
 
 def _make_sequences() -> _np_cpu.ndarray:
-    """Binary sequences: each row is one sequence of T frames."""
     rng  = _np_cpu.random.RandomState(0)
     base = (rng.rand(NUM_SEQ, SENSORY_SIZE) > 0.5).astype(_np_cpu.float64)
-    seqs = _np_cpu.stack([base] * T, axis=1)   # (NUM_SEQ, T, SENSORY_SIZE)
+    return _np_cpu.stack([base] * T, axis=1)   # (NUM_SEQ, T, SENSORY_SIZE)
-    return seqs
 
 
-def test_rnn_single_layer():
+def test_rnn_shared():
+    """Shared-weights mode: one entity reused at every time step."""
     seqs = _make_sequences()
+    rnn  = StackRnn("test_rnn_shared", WORK_DIR)
+    rnn.append(StackRnn.make_layer("layer0", SENSORY_SIZE, H_SIZE,
+                                   EntityParams(), _PARAMS))
+    rnn.state_init(0.01)
 
-    rnn = StackRnn("test_rnn", WORK_DIR)
+    assert rnn.is_shared
-    layer = StackRnn.make_layer(
+    assert rnn.sensory_size() == SENSORY_SIZE
-        "layer0", SENSORY_SIZE, H_SIZE,
+    assert rnn.h_size()       == H_SIZE
-        EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False),
+
-        TrainingParams(learning_rate=0.05, momentum=0.5, num_epochs=20,
+    rnn.train(seqs)
-                       mini_batch_size=0, do_rao_blackwell=True),
+
-    )
+    rnn.reset(batch_size=1)
+    for t in range(T):
+        h = rnn.step(np.array(seqs[0, t][None, :]))
+        assert h.shape == (1, H_SIZE), f"bad shape at t={t}"
+
+    recon = rnn.reconstruct(h)
+    assert recon.shape == (1, SENSORY_SIZE)
+
+    print(f"Original : {convert(np.array(seqs[0, -1][None, :]))}")
+    print(f"Recon    : {convert(recon)}")
+    print("test_rnn_shared: [passed]")
+
+
+def test_rnn_unrolled():
+    """Unrolled mode: T entities, one per time step, each with own weights."""
+    seqs = _make_sequences()
+    rnn  = StackRnn("test_rnn_unrolled", WORK_DIR)
+    for layer in StackRnn.make_unrolled(T, SENSORY_SIZE, H_SIZE,
+                                        EntityParams(), _PARAMS):
         rnn.append(layer)
     rnn.state_init(0.01)
 
-    # Confirm entity shape
+    assert not rnn.is_shared
-    assert rnn.sensory_size() == SENSORY_SIZE, "sensory_size mismatch"
+    assert rnn.num_layers() == T
-    assert rnn.h_size() == H_SIZE, "h_size mismatch"
 
     rnn.train(seqs)
 
-    # Inference: step through one sequence
+    # Inference — _t advances through all T entities
     rnn.reset(batch_size=1)
-    seq0 = seqs[0]                             # (T, SENSORY_SIZE) numpy
     for t in range(T):
-        x_t = np.array(seq0[t][None, :])      # (1, SENSORY_SIZE) on device
+        h = rnn.step(np.array(seqs[0, t][None, :]))
-        h = rnn.step(x_t)
+        assert h.shape == (1, H_SIZE), f"bad shape at t={t}"
-        assert h.shape == (1, H_SIZE), f"step output shape wrong at t={t}"
+        assert rnn._t == t + 1
 
-    # Reconstruction from final hidden state
     recon = rnn.reconstruct(h)
-    assert recon.shape == (1, SENSORY_SIZE), "reconstruct shape wrong"
+    assert recon.shape == (1, SENSORY_SIZE)
 
-    print(f"Original x_T : {convert(np.array(seq0[-1][None, :]))}")
+    # next_entity wraps around modularly
-    print(f"Reconstructed: {convert(recon)}")
+    rnn._t = T + 3
-    print("test_rnn_single_layer: [passed]")
+    assert rnn.next_entity() is rnn.from_index(3).entity
 
-
+    print("test_rnn_unrolled: [passed]")
-def test_rnn_two_layers():
-    """Two-layer recurrent stack: layer 1 receives h_0 as its sensory input."""
-    H_SIZE_0, H_SIZE_1 = 6, 3
-    seqs = _make_sequences()
-
-    rnn = StackRnn("test_rnn2", WORK_DIR)
-    rnn.append(StackRnn.make_layer(
-        "layer0", SENSORY_SIZE, H_SIZE_0,
-        EntityParams(),
-        TrainingParams(learning_rate=0.05, num_epochs=10),
-    ))
-    rnn.append(StackRnn.make_layer(
-        "layer1", H_SIZE_0, H_SIZE_1,
-        EntityParams(),
-        TrainingParams(learning_rate=0.05, num_epochs=10),
-    ))
-    rnn.state_init(0.01)
-
-    rnn.train(seqs)
-
-    rnn.reset(batch_size=1)
-    for t in range(T):
-        x_t = np.array(seqs[0, t][None, :])
-        h = rnn.step(x_t)
-
-    assert h.shape == (1, H_SIZE_1)
-    print("test_rnn_two_layers: [passed]")
 
 
 def test_rnn_save_load():
     seqs = _make_sequences()
-
     rnn  = StackRnn("test_rnn_sl", WORK_DIR)
-    rnn.append(StackRnn.make_layer(
+    for layer in StackRnn.make_unrolled(T, SENSORY_SIZE, H_SIZE,
-        "layer0", SENSORY_SIZE, H_SIZE,
+                                        EntityParams(), _PARAMS):
-        EntityParams(),
+        rnn.append(layer)
-        TrainingParams(num_epochs=5),
-    ))
     rnn.state_init(0.01)
     rnn.train(seqs)
     rnn.state_save()
@@ -106,7 +93,7 @@ def test_rnn_save_load():
 
 
 if __name__ == "__main__":
-    test_rnn_single_layer()
+    test_rnn_shared()
-    test_rnn_two_layers()
+    test_rnn_unrolled()
     test_rnn_save_load()
     print("All RNN tests passed.")