Publish frozen LoopFormer pointer checkpoint and inference code

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,80 @@
+---
+language: en
+base_model: Qwen/Qwen2.5-0.5B-Instruct
+tags:
+  - loopformer
+  - recurrent-transformer
+  - pointer-chasing
+  - custom-code
+---
+
+# LoopFormer: recurrent pointer execution
+
+Frozen final checkpoint: `controller-shared-number-precision-seed83`.
+This is a custom recurrent Qwen model for following a supplied A–Z rule table.
+It requires the included Python inference code; ordinary `AutoModelForCausalLM`
+generation does not implement its recurrence.
+
+## Download and run
+
+Use Python 3.11 or 3.12. Install a PyTorch build suitable for your device.
+The requirements record the validated research dependency versions; the Space
+adds Gradio and `spaces` for hosting. Export a graph request from the playground,
+or create JSON with `mapping` (26 ordered `[source, destination]` pairs), `start`
+(one A–Z letter), and `steps` (a positive integer).
+
+```bash
+hf download mukunds/loopformer --local-dir loopformer-model
+cd loopformer-model
+python -m pip install -r requirements.txt
+python -m scripts.demo --checkpoint . --request graph.json --device cuda --download
+```
+
+The explicit `--download` option permits downloading the pinned Qwen base
+revision. It is needed on the first run. `--device cpu` and `--device mps` are
+also available, but their latency and final-checkpoint fidelity must be measured
+on the selected device. Each output line is a cumulative JSON snapshot of actual
+model predictions; the final line includes separate correctness/stopping metrics.
+
+The default requested-step ceiling is 256 and the safety cap is 272. To explore
+longer rollouts, explicitly increase both `--max-steps` and `--loop-cap`. This
+does not establish graph-execution accuracy beyond the evaluated 256 steps.
+The symbol vocabulary remains exactly A–Z; 256 steps does not mean 256 nodes.
+
+## Architecture and inference
+
+The base is `Qwen/Qwen2.5-0.5B-Instruct`, revision
+`7ae557604adf67be50417f59c2c2f167def9a775`. Six prelude layers prepare the prompt;
+twelve shared middle layers update a recurrent working state; six coda layers
+read out a distribution over the 26 letters after every loop. A trained bridge
+re-enters the shared middle block. Fixed prefix memory is reused within a run.
+The final executor adapts all R weights and freezes P/C.
+
+The Steps field is routed away from the executor. A shared learned number reader
+initializes a separate scalar timer, whose learned affine recurrence and stop
+readout determine completion. Stop probability is a timing signal, not a
+correctness estimate. Decoded predictions and reference states are never fed
+back into recurrence. Default inference is FP32 with SDPA and threshold 0.5.
+
+`adapter_model.pt` contains all adapted executor/controller tensors, not merely
+LoRA weights. SHA-256:
+`0f8a3f2572ec2caf4fafbcdd074b5283bd8c610602502763b3c04487ce1f123d`.
+Optimizer state and training data are excluded from this inference release.
+
+## Evidence and limitations
+
+On the frozen final panel of 1,350 new graphs and requests 1–256:
+
+- Final letter plus exact stopping: 97.81% over 345,600 graph/count queries.
+- Complete trajectories through depth 256: 96.89% (1,308/1,350 graphs).
+- Exact stopping: 100% across that finite requested-count panel.
+
+Executor supervision used depths 1–6, 8, 10, 12. Controller fitting used
+counts 1–63 excluding 9, 17, 29, 41, 53, with supervised unrolls capped at 12.
+Controller-only numerical results at longer horizons are distinct from executing
+graphs correctly. This is one task family, not general reasoning or transfer
+evidence. Arbitrary hand-edited graphs can fail. Changes to precision, runtime,
+hardware or kernels need validation against native intermediate readouts.
+
+See the [source and full evidence](https://github.com/MukundShankar-dev/loopformer)
+and the [research journal](https://mukundshankar-dev.github.io/projects/loopformer/).

adapter_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f8a3f2572ec2caf4fafbcdd074b5283bd8c610602502763b3c04487ce1f123d
+size 719156751

chat_template.jinja

@@ -0,0 +1,54 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] }}
+    {%- else %}
+        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}
+    {%- endif %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
+    {%- else %}
+        {{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role }}
+        {%- if message.content %}
+            {{- '\n' + message.content }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call.function is defined %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {{- '\n<tool_call>\n{"name": "' }}
+            {{- tool_call.name }}
+            {{- '", "arguments": ' }}
+            {{- tool_call.arguments | tojson }}
+            {{- '}\n</tool_call>' }}
+        {%- endfor %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}

recurrent_config.json

@@ -0,0 +1,241 @@
+{
+  "base_model": "Qwen/Qwen2.5-0.5B-Instruct",
+  "revision": "7ae557604adf67be50417f59c2c2f167def9a775",
+  "recurrent_start": 6,
+  "recurrent_end": 18,
+  "lora_rank": 8,
+  "lora_alpha": 16,
+  "symbols": [
+    "A",
+    "B",
+    "C",
+    "D",
+    "E",
+    "F",
+    "G",
+    "H",
+    "I",
+    "J",
+    "K",
+    "L",
+    "M",
+    "N",
+    "O",
+    "P",
+    "Q",
+    "R",
+    "S",
+    "T",
+    "U",
+    "V",
+    "W",
+    "X",
+    "Y",
+    "Z"
+  ],
+  "token_ids": [
+    362,
+    425,
+    356,
+    422,
+    468,
+    434,
+    479,
+    472,
+    358,
+    619,
+    730,
+    444,
+    386,
+    451,
+    506,
+    393,
+    1207,
+    431,
+    328,
+    350,
+    547,
+    647,
+    467,
+    1599,
+    809,
+    1863
+  ],
+  "prompt_format": "dataset_raw",
+  "loss_vocabulary": "symbols",
+  "train_max_depth": 12,
+  "train_depths": [
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    8,
+    10,
+    12
+  ],
+  "recurrence_mode": "fixed_prompt",
+  "completion_head": {
+    "intermediate": 1
+  },
+  "completion_loss_weight": 0.1,
+  "training_compute": {
+    "precision": "bf16",
+    "attention": "sdpa"
+  },
+  "executor": {
+    "train_scope": "recurrent",
+    "isolated": true,
+    "bridge": true,
+    "direct_readout": true,
+    "controller_size": 1,
+    "prefix_reuse": true,
+    "gradient_checkpointing": true,
+    "controller_kind": "shared_number"
+  },
+  "format": "loopformer-executor-v2",
+  "controller_training": {
+    "source": "models/stage1_pointer/controller-shared-number-seed83/best",
+    "source_checkpoint_sha256": "42059fd6fd4cc1ffe05a705a99a7aba76d4e7e0d87933bad18c7adc4cdee1110",
+    "trained_depths": [
+      1,
+      2,
+      3,
+      4,
+      5,
+      6,
+      7,
+      8,
+      10,
+      11,
+      12,
+      13,
+      14,
+      15,
+      16,
+      18,
+      19,
+      20,
+      21,
+      22,
+      23,
+      24,
+      25,
+      26,
+      27,
+      28,
+      30,
+      31,
+      32,
+      33,
+      34,
+      35,
+      36,
+      37,
+      38,
+      39,
+      40,
+      42,
+      43,
+      44,
+      45,
+      46,
+      47,
+      48,
+      49,
+      50,
+      51,
+      52,
+      54,
+      55,
+      56,
+      57,
+      58,
+      59,
+      60,
+      61,
+      62,
+      63
+    ],
+    "training_loops": 12,
+    "optimization_scope": "countdown_cell_only",
+    "config": {
+      "source": "models/stage1_pointer/controller-shared-number-seed83/best",
+      "training_config": "configs/controller_affine.json",
+      "output": "models/stage1_pointer/controller-shared-number-precision-seed83"
+    },
+    "fit": {
+      "method": "float64 L-BFGS on original free-running numerical prefix labels; float32 deployment",
+      "training_counts": [
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        7,
+        8,
+        10,
+        11,
+        12,
+        13,
+        14,
+        15,
+        16,
+        18,
+        19,
+        20,
+        21,
+        22,
+        23,
+        24,
+        25,
+        26,
+        27,
+        28,
+        30,
+        31,
+        32,
+        33,
+        34,
+        35,
+        36,
+        37,
+        38,
+        39,
+        40,
+        42,
+        43,
+        44,
+        45,
+        46,
+        47,
+        48,
+        49,
+        50,
+        51,
+        52,
+        54,
+        55,
+        56,
+        57,
+        58,
+        59,
+        60,
+        61,
+        62,
+        63
+      ],
+      "training_loops": 12,
+      "objective_evaluations": 10,
+      "before_mse": 4.1626904266916364e-09,
+      "fitted_mse": 2.002921115782509e-12,
+      "deployed_mse": 6.307695082326698e-12,
+      "learned_gain": 1.0,
+      "learned_offset": -0.9999999403953552,
+      "numeric_targets_are_inputs": false
+    },
+    "selection": "One predetermined training-label fit; no stress/graph benchmark selection"
+  }
+}

requirements.txt

@@ -0,0 +1,6 @@
+torch==2.14.0
+transformers==5.17.0
+tokenizers==0.23.2
+huggingface-hub==1.31.0
+peft==0.20.0
+numpy==2.4.6

scripts/dataset/__init__.py

@@ -0,0 +1 @@
+"""Synthetic tasks with exact reference execution; no model loading."""

scripts/dataset/__main__.py

@@ -0,0 +1,7 @@
+"""Run the dataset CLI with python -m scripts.dataset."""
+
+from scripts.dataset.cli import main
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset/benchmark.py

@@ -0,0 +1,57 @@
+"""Independent, seed-replayable graph panels for frozen pointer benchmarks."""
+from dataclasses import replace
+import json
+from pathlib import Path
+
+from .pointer import (PointerExample, example_seed, execute, generate_unconditioned_example,
+                      orbit_structure, render_prompt, validate_example)
+
+
+def count_variant(task: PointerExample, count: int) -> PointerExample:
+    """Change requested horizon/labels only; never condition graph sampling on it."""
+    states = execute(dict(task.mapping), task.initial_state, count)
+    result = replace(task, example_id=f'{task.example_id}-steps-{count}', task_depth=count,
+                     intermediate_states=states, final_state=states[-1],
+                     prompt=render_prompt(task.mapping, task.initial_state, count))
+    validate_example(result)
+    return result
+
+
+def independent_graphs(config: dict, excluded: set[str]) -> tuple[list[PointerExample], list[dict]]:
+    """Balanced seed/mode strata; reject only graph identity overlap, never difficulty."""
+    if config['graphs_per_seed_mode'] < 1 or not 1 <= config['max_depth'] <= 256:
+        raise ValueError('Need positive graph count and a maximum depth in 1..256')
+    if len(set(config['seeds'])) != len(config['seeds']) or len(set(config['graph_modes'])) != len(config['graph_modes']):
+        raise ValueError('Repeated seeds or graph modes')
+    used = set(excluded)
+    tasks, metadata = [], []
+    for seed in config['seeds']:
+        for mode in config['graph_modes']:
+            for index in range(config['graphs_per_seed_mode']):
+                for attempt in range(1000):
+                    task = generate_unconditioned_example(example_seed(seed, f'benchmark-{mode}', index, attempt),
+                        config['max_depth'], 'benchmark', index, mode)
+                    if task.mapping_sha256 not in used:
+                        break
+                else:
+                    raise ValueError('Unable to sample a new rule table after 1000 attempts')
+                used.add(task.mapping_sha256)
+                task = replace(task, example_id=f'benchmark-{seed}-{mode}-{index:06d}-{task.seed:016x}')
+                validate_example(task)
+                tasks.append(task)
+                metadata.append({'graph_index': len(tasks) - 1, 'example_id': task.example_id,
+                    'mapping_sha256': task.mapping_sha256, 'dataset_seed': seed, 'graph_mode': mode,
+                    **orbit_structure(dict(task.mapping), task.initial_state)})
+    return tasks, metadata
+
+
+def exclusion_hashes(directory: Path) -> set[str]:
+    """Read graph fingerprints, including reserved splits, solely to prevent overlap."""
+    hashes = set()
+    for path in sorted(directory.glob('*.jsonl')):
+        for line in path.open():
+            record = json.loads(line)
+            hashes.add(record['mapping_sha256'])
+    if not hashes:
+        raise ValueError('No old dataset graph fingerprints found')
+    return hashes

scripts/dataset/cli.py

@@ -0,0 +1,128 @@
+"""Generate or verify reproducible Stage 1 pointer data (no model training).
+
+Run from the repository root: python -m scripts.dataset --help
+"""
+
+import argparse
+from dataclasses import asdict
+import hashlib
+from importlib.metadata import version
+from pathlib import Path
+import platform
+import subprocess
+import sys
+
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+
+from scripts.dataset.dataset import (
+    DatasetConfig, SPLITS, generate_dataset, make_record, preview_indices,
+    verify_dataset, write_dataset,
+)
+from scripts.dataset.symbols import MODEL_ID, MODEL_REVISION, validate_symbols
+
+
+def show_summary(console: Console, manifest: dict) -> None:
+    table = Table(title="Validated pointer dataset")
+    for column in ("Split", "Examples", "Depths (count)", "Prompt tokens"):
+        table.add_column(column)
+    for split in SPLITS:
+        summary = manifest["splits"][split]
+        depths = ", ".join(f"{depth}: {count}" for depth, count in sorted(summary["depth_counts"].items(), key=lambda item: int(item[0])))
+        table.add_row(split, f"{summary['count']:,}", depths, f"{summary['prompt_tokens_min']}–{summary['prompt_tokens_max']}")
+    console.print(table)
+    console.print("[green]PASS[/green] Exact targets, tokenizer contexts, declared graph sharing, disjoint split boundaries, and seed replay")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--graph-mode", choices=("unique_path", "mixture", "random_function", "permutation", "full_cycle"), default="unique_path")
+    parser.add_argument("--paired-horizons", action="store_true", help="Match each evaluation graph/start across depths; training graphs stay independent")
+    parser.add_argument("--seed", type=int, default=17)
+    parser.add_argument("--train-count", type=int, default=10_000)
+    parser.add_argument("--validation-count", type=int, default=1_000)
+    parser.add_argument("--test-count", type=int, default=1_000)
+    parser.add_argument("--depth-test-count", type=int, default=1_000)
+    parser.add_argument("--min-depth", type=int, default=1)
+    parser.add_argument("--max-train-depth", type=int, default=8)
+    parser.add_argument("--max-eval-depth", type=int, default=16)
+    parser.add_argument("--train-depths", type=int, nargs="+", help="Explicit training depths; validation/test retain every depth through max-train-depth")
+    parser.add_argument("--output", type=Path, help="New directory; default data/pointer/seed-<seed>")
+    mode = parser.add_mutually_exclusive_group()
+    mode.add_argument("--dry-run", nargs="?", const=5, type=int, choices=(5, 10), help="Preview up to 5 (default) or 10 examples spanning depths; write no dataset files")
+    mode.add_argument("--verify", type=Path, metavar="DIRECTORY", help="Read and independently verify an existing dataset")
+    args = parser.parse_args()
+    config = DatasetConfig(**{key: getattr(args, key) for key in asdict(DatasetConfig())})
+    try:
+        config.validate()
+    except ValueError as error:
+        parser.error(str(error))
+    output = args.output or Path("data/pointer") / f"seed-{args.seed}"
+    if args.verify is None and args.dry_run is None and output.exists():
+        parser.error(f"Output already exists: {output}; choose a new --output (never overwritten)")
+
+    from transformers import AutoTokenizer
+
+    console = Console()
+    console.print("[bold cyan]Stage 1 · pointer data[/bold cyan]")
+    console.print(f"Cached tokenizer only: {MODEL_ID} @ {MODEL_REVISION}")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=MODEL_REVISION, local_files_only=True)
+    if args.verify is not None:
+        with console.status("Verifying saved files and replaying seeds…"):
+            manifest = verify_dataset(args.verify, tokenizer)
+        show_summary(console, manifest)
+        console.print(f"Verified: {args.verify.resolve()}")
+        return
+
+    token_ids = validate_symbols(tokenizer)
+    console.print(f"[green]PASS[/green] 26 unique single-token symbols in rule, start, and answer contexts; seed {config.seed}")
+    console.print(f"Training depths: {list(config.depths('train'))}; validation/test: {config.min_depth}–{config.max_train_depth}; extrapolation depths: {config.max_train_depth + 1}–{config.max_eval_depth}")
+    if args.dry_run is not None:
+        distribution = Table(title="Configured dataset distribution (not preview frequencies)")
+        for column in ("Split", "Total", "Steps: example count"):
+            distribution.add_column(column)
+        for split in SPLITS:
+            depths = config.depths(split)
+            quotient, remainder = divmod(config.count(split), len(depths))
+            counts = ", ".join(f"{depth}: {quotient + (index < remainder)}" for index, depth in enumerate(depths))
+            distribution.add_row(split, f"{config.count(split):,}", counts)
+        console.print(distribution)
+        selected = preview_indices(config, args.dry_run)
+        console.print("[cyan]Preview spans shallow through deep tasks; it is not a random frequency sample.[/cyan]")
+        for split, index in selected:
+            record = make_record(config, split, index, tokenizer, token_ids)
+            console.print(Panel(Text(record["prompt"]), title=f"{record['example_id']} · depth {record['task_depth']}"))
+            table = Table(title=f"Start {record['initial_state']} · seed {record['seed']}")
+            for column in ("Step / loop", "Expected symbol", "Answer token ID"):
+                table.add_column(column)
+            for step, (symbol, token_id) in enumerate(zip(record["intermediate_states"], record["intermediate_token_ids"]), 1):
+                table.add_row(str(step), symbol, str(token_id))
+            console.print(table)
+            console.print(f"[bold green]Final: {record['final_state']}[/bold green] · answer position {record['answer_position']}")
+        console.print(f"[yellow]Dry run complete: {len(selected)} examples; no dataset files written.[/yellow]")
+        return
+
+    root = Path(__file__).resolve().parents[2]
+    sources = sorted((root / "scripts/dataset").glob("*.py"))
+    provenance = {
+        "command": [sys.executable, "-m", "scripts.dataset", *sys.argv[1:]],
+        "python": platform.python_version(),
+        "packages": {name: version(name) for name in ("transformers", "tokenizers", "huggingface-hub", "rich")},
+        "git_head": subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=root, text=True).strip(),
+        "git_status": subprocess.check_output(["git", "status", "--short"], cwd=root, text=True),
+        "source_sha256": {str(path.relative_to(root)): hashlib.sha256(path.read_bytes()).hexdigest() for path in sources},
+    }
+    with console.status("Generating and checking every example…"):
+        dataset = generate_dataset(config, tokenizer, token_ids)
+        write_dataset(output, dataset, config, token_ids, provenance)
+    with console.status("Reading files back and independently replaying seeds…"):
+        manifest = verify_dataset(output, tokenizer)
+    show_summary(console, manifest)
+    console.print(f"Dataset and manifest: [bold]{output.resolve()}[/bold]")
+    console.print("Data milestone complete. No training performed.")
+
+
+if __name__ == "__main__":
+    main()

scripts/dataset/dataset.py

@@ -0,0 +1,212 @@
+"""Seeded splits, JSONL persistence, and independent pointer-data validation."""
+
+from collections import Counter
+from dataclasses import asdict, dataclass
+import hashlib
+import json
+from pathlib import Path
+from typing import Any
+
+from scripts.dataset.pointer import (PointerExample, example_seed, generate_example, validate_example,
+                                     generate_unconditioned_example, orbit_structure)
+from scripts.dataset.symbols import MODEL_ID, MODEL_REVISION, validate_prompt_tokens, validate_symbols
+
+
+SPLITS = ("train", "validation", "test", "depth_test")
+
+
+@dataclass(frozen=True)
+class DatasetConfig:
+    graph_mode: str = "unique_path"
+    paired_horizons: bool = False
+    seed: int = 17
+    train_count: int = 10_000
+    validation_count: int = 1_000
+    test_count: int = 1_000
+    depth_test_count: int = 1_000
+    min_depth: int = 1
+    max_train_depth: int = 8
+    max_eval_depth: int = 16
+    train_depths: list[int] | None = None
+
+    def validate(self) -> None:
+        if any(type(value) is not int for name, value in asdict(self).items() if name not in ("train_depths", "graph_mode", "paired_horizons")):
+            raise ValueError("Dataset configuration values must be integers")
+        if self.graph_mode not in ("unique_path", "mixture", "random_function", "permutation", "full_cycle"):
+            raise ValueError("Unknown graph_mode")
+        if type(self.paired_horizons) is not bool or (self.paired_horizons and self.graph_mode == "unique_path"):
+            raise ValueError("paired_horizons requires depth-independent graph sampling")
+        if self.seed < 0 or min(self.count(split) for split in SPLITS) < 1:
+            raise ValueError("seed must be nonnegative and all split counts must be positive")
+        if not 1 <= self.min_depth <= self.max_train_depth < self.max_eval_depth <= (25 if self.graph_mode == "unique_path" else 256):
+            raise ValueError("Require ordered positive train/eval depths within the graph-mode limit")
+
+        if self.train_depths is not None:
+            if (not self.train_depths or any(type(d) is not int for d in self.train_depths)
+                    or sorted(set(self.train_depths)) != self.train_depths
+                    or min(self.train_depths) < self.min_depth or max(self.train_depths) != self.max_train_depth):
+                raise ValueError("train_depths must be sorted unique integers within bounds, ending at max_train_depth")
+
+    def paired(self, split: str) -> bool:
+        """Match evaluation horizons; retain independent graphs in training."""
+        return self.paired_horizons and split != "train"
+
+    def count(self, split: str) -> int:
+        if split not in SPLITS:
+            raise ValueError(f"Unknown split: {split}")
+        return getattr(self, f"{split}_count")
+
+    def depths(self, split: str) -> range | list[int]:
+        if split not in SPLITS:
+            raise ValueError(f"Unknown split: {split}")
+        if split == "train" and self.train_depths is not None:
+            return self.train_depths
+        if split == "depth_test":
+            return range(self.max_train_depth + 1, self.max_eval_depth + 1)
+        return range(self.min_depth, self.max_train_depth + 1)
+
+
+def preview_indices(config: DatasetConfig, count: int) -> list[tuple[str, int]]:
+    """Select actual records spanning available depths, without sampling programs.
+
+    This is a depth showcase, not a random sample of split frequencies. Keep only
+    enough candidate indices to fill a preview, even for very large split counts.
+    """
+    config.validate()
+    if count not in (5, 10):
+        raise ValueError("Preview count must be 5 or 10")
+    candidates = [
+        (split, index, config.depths(split)[index % len(config.depths(split))])
+        for split in SPLITS
+        for index in range(min(config.count(split), count * len(config.depths(split))))
+    ]
+    depths = sorted({depth for _, _, depth in candidates})
+    count = min(count, len(candidates))
+    selected = []
+    for position in range(count):
+        target = depths[position * (len(depths) - 1) // (count - 1)]
+        candidate = min(candidates, key=lambda item: (
+            abs(item[2] - target), (SPLITS.index(item[0]) - position) % len(SPLITS), item[1],
+        ))
+        candidates.remove(candidate)
+        selected.append(candidate[:2])
+    return selected
+
+
+def make_record(config: DatasetConfig, split: str, index: int, tokenizer: Any, token_ids: dict[str, int]) -> dict:
+    """Index determines balanced depth and an independently derived random seed."""
+    depths = config.depths(split)
+    depth = depths[index % len(depths)]
+    graph_index = index // len(depths) if config.paired(split) else index
+    seed = example_seed(config.seed, split, graph_index)
+    example = (generate_example(seed, depth, split, index) if config.graph_mode == "unique_path" else
+               generate_unconditioned_example(seed, depth, split, index, config.graph_mode))
+    validate_example(example)
+    position = validate_prompt_tokens(tokenizer, example.prompt, token_ids)
+    return {
+        **example.to_dict(),
+        **({"graph_sampling": config.graph_mode, **orbit_structure(dict(example.mapping), example.initial_state)}
+           if config.graph_mode != "unique_path" else {}),
+        "answer_position": position,
+        "prompt_token_count": position + 1,
+        "intermediate_token_ids": [token_ids[symbol] for symbol in example.intermediate_states],
+        "final_token_id": token_ids[example.final_state],
+    }
+
+
+def generate_dataset(config: DatasetConfig, tokenizer: Any, token_ids: dict[str, int]) -> dict[str, list[dict]]:
+    config.validate()
+    dataset = {}
+    seen = {}
+    for split in SPLITS:
+        records = []
+        for index in range(config.count(split)):
+            record = make_record(config, split, index, tokenizer, token_ids)
+            fingerprint = record["mapping_sha256"]
+            if fingerprint in seen and (not config.paired(split) or seen[fingerprint] != split):
+                raise ValueError("Repeated rule table within/across splits; select a different master seed")
+            seen[fingerprint] = split
+            records.append(record)
+        dataset[split] = records
+    return dataset
+
+
+def encode_records(records: list[dict]) -> bytes:
+    return "".join(json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n" for record in records).encode()
+
+
+def summarize(records: list[dict]) -> dict:
+    return {
+        "count": len(records),
+        "depth_counts": dict(sorted(Counter(str(record["task_depth"]) for record in records).items())),
+        "prompt_tokens_min": min(record["prompt_token_count"] for record in records),
+        "prompt_tokens_max": max(record["prompt_token_count"] for record in records),
+        "sha256": hashlib.sha256(encode_records(records)).hexdigest(),
+    }
+
+
+def write_dataset(
+    output: Path, dataset: dict[str, list[dict]], config: DatasetConfig,
+    token_ids: dict[str, int], provenance: dict,
+) -> dict:
+    """Write a new directory only; the manifest is the final completion marker."""
+    manifest = {
+        "schema_version": 1,
+        "generator": "pointer-v1" if config.graph_mode == "unique_path" else "pointer-v2",
+        "config": asdict(config),
+        "tokenizer": {
+            "model_id": MODEL_ID, "revision": MODEL_REVISION,
+            "add_special_tokens": False, "chat_template": False,
+            "answer_text_prefix": " ", "symbol_token_ids": token_ids,
+        },
+        "semantics": ("explicit_steps_unique_nominal_path_no_post_completion_targets" if config.graph_mode == "unique_path"
+                      else "depth_independent_graphs_exact_steps_cycles_allowed"),
+        "splits": {split: summarize(dataset[split]) for split in SPLITS},
+        "validation": {"unique_mapping_tables": not config.paired_horizons, "exact_targets": True, "prompt_symbol_tokens": True},
+        "provenance": provenance,
+    }
+    output.mkdir(parents=True, exist_ok=False)
+    for split in SPLITS:
+        (output / f"{split}.jsonl").write_bytes(encode_records(dataset[split]))
+    (output / "manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
+    return manifest
+
+
+def verify_dataset(output: Path, tokenizer: Any) -> dict:
+    """Check bytes, labels, split boundaries, tokenizer contexts, and seed replay.
+
+    Regeneration alone is insufficient: validate_example separately reparses the
+    rendered rules and executes them to check the actual stored trajectory.
+    """
+    manifest = json.loads((output / "manifest.json").read_text())
+    if manifest["schema_version"] != 1 or manifest["generator"] not in ("pointer-v1", "pointer-v2"):
+        raise ValueError("Unsupported dataset version")
+    config = DatasetConfig(**manifest["config"])
+    config.validate()
+    expected_generator = "pointer-v1" if config.graph_mode == "unique_path" else "pointer-v2"
+    if manifest["generator"] != expected_generator:
+        raise ValueError("Generator version and graph sampling mode disagree")
+    token_ids = validate_symbols(tokenizer)
+    if manifest["tokenizer"] != {
+        "model_id": MODEL_ID, "revision": MODEL_REVISION,
+        "add_special_tokens": False, "chat_template": False,
+        "answer_text_prefix": " ", "symbol_token_ids": token_ids,
+    }:
+        raise ValueError("Tokenizer metadata mismatch")
+    seen = {}
+    for split in SPLITS:
+        payload = (output / f"{split}.jsonl").read_bytes()
+        if hashlib.sha256(payload).hexdigest() != manifest["splits"][split]["sha256"]:
+            raise ValueError(f"{split}: file checksum mismatch")
+        records = [json.loads(line) for line in payload.splitlines()]
+        if len(records) != config.count(split) or summarize(records) != manifest["splits"][split]:
+            raise ValueError(f"{split}: split count/summary mismatch")
+        for index, record in enumerate(records):
+            example = PointerExample(**{key: record[key] for key in PointerExample.__dataclass_fields__})
+            validate_example(example)
+            if record != make_record(config, split, index, tokenizer, token_ids):
+                raise ValueError(f"{split} record {index}: seed replay/token metadata mismatch")
+            if example.mapping_sha256 in seen and (not config.paired(split) or seen[example.mapping_sha256] != split):
+                raise ValueError("Repeated rule table within/across splits")
+            seen[example.mapping_sha256] = split
+    return manifest

scripts/dataset/pointer.py

@@ -0,0 +1,189 @@
+"""Explicit-step pointer programs and exact, one-based execution targets."""
+
+from dataclasses import asdict, dataclass
+import hashlib
+import json
+import random
+import re
+from typing import Sequence
+
+
+SYMBOLS = tuple("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
+PAIR = re.compile(r"\(\s*([A-Z])\s*,\s*([A-Z])\s*\)")
+
+
+def parse_mapping(text: str) -> dict[str, str]:
+    """Parse pairs such as '(A,C) (C,D) (D, E)', rejecting ambiguous rules."""
+    matches = list(PAIR.finditer(text))
+    if not matches or PAIR.sub("", text).strip():
+        raise ValueError("Mapping must be whitespace-separated (A,C) pairs")
+    mapping = {}
+    for match in matches:
+        source, target = match.groups()
+        if source in mapping:
+            raise ValueError(f"Duplicate source symbol: {source}")
+        mapping[source] = target
+    return mapping
+
+
+def execute(mapping: dict[str, str], start: str, steps: int) -> list[str]:
+    """Return states after transitions 1..steps, excluding the initial state.
+
+    This reference interpreter permits cycles and partial tables, provided every
+    requested transition exists. The dataset sampler applies stricter constraints.
+    """
+    if type(steps) is not int or steps < 0:
+        raise ValueError("steps must be a nonnegative integer")
+    states = []
+    current = start
+    for step in range(1, steps + 1):
+        if current not in mapping:
+            raise ValueError(f"No rule for {current!r} at step {step}")
+        current = mapping[current]
+        states.append(current)
+    return states
+
+
+def check_predictions(
+    mapping: dict[str, str], start: str, predictions: Sequence[str], steps: int,
+) -> list[bool]:
+    """Score decoded symbols at steps 1..N against the exact nominal trajectory.
+
+    Predictions may be a prefix (too few model loops), but cannot extend beyond
+    the requested depth. Invalid/multi-symbol answers are incorrect; no lenient
+    extraction of a letter from prose is performed. Leading whitespace is allowed.
+    """
+    targets = execute(mapping, start, steps)
+    if isinstance(predictions, str) or len(predictions) > steps:
+        raise ValueError("Provide a sequence of at most 'steps' decoded predictions")
+    return [prediction.strip() == target for prediction, target in zip(predictions, targets)]
+
+
+def render_prompt(pairs: list[list[str]], start: str, steps: int) -> str:
+    # Spaces before symbols prevent Qwen punctuation/letter token merges.
+    rules = " ".join(f"( {source}, {target})" for source, target in pairs)
+    return f"Rules: {rules}\nStart: {start}\nSteps: {steps}\nAnswer:"
+
+
+def mapping_fingerprint(mapping: dict[str, str]) -> str:
+    """Ignore rule order, start, and depth when detecting reused rule tables."""
+    payload = json.dumps(sorted(mapping.items()), separators=(",", ":"))
+    return hashlib.sha256(payload.encode()).hexdigest()
+
+
+def example_seed(master_seed: int, split: str, index: int, attempt: int = 0) -> int:
+    payload = f"pointer-v1:{master_seed}:{split}:{index}:{attempt}"
+    return int.from_bytes(hashlib.sha256(payload.encode()).digest()[:8], "big")
+
+
+@dataclass
+class PointerExample:
+    schema_version: int
+    example_id: str
+    family: str
+    split: str
+    seed: int
+    task_depth: int
+    mapping: list[list[str]]
+    mapping_sha256: str
+    initial_state: str
+    intermediate_states: list[str]
+    final_state: str
+    prompt: str
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+def generate_example(seed: int, depth: int, split: str, index: int) -> PointerExample:
+    """Sample a full table, conditioned on a non-repeating nominal path.
+
+    Every example uses all 26 symbols, keeping rule count independent of depth.
+    The d+1 distinct path states prevent an early/repeated final answer; remaining
+    edges are independent random choices and can cycle outside the nominal path.
+    Rule order is shuffled so the solution is not printed in execution order.
+    """
+    if type(depth) is not int or not 1 <= depth < len(SYMBOLS):
+        raise ValueError("depth must be an integer between 1 and 25")
+    if type(seed) is not int or seed < 0:
+        raise ValueError("seed must be a nonnegative integer")
+    rng = random.Random(seed)
+    path = rng.sample(SYMBOLS, depth + 1)
+    mapping = dict(zip(path[:-1], path[1:]))
+    for symbol in SYMBOLS:
+        if symbol not in mapping:
+            mapping[symbol] = rng.choice(SYMBOLS)
+    pairs = [[source, mapping[source]] for source in SYMBOLS]
+    rng.shuffle(pairs)
+    states = execute(mapping, path[0], depth)
+    return PointerExample(
+        schema_version=1, example_id=f"{split}-{index:06d}-{seed:016x}",
+        family="pointer", split=split, seed=seed, task_depth=depth,
+        mapping=pairs, mapping_sha256=mapping_fingerprint(mapping),
+        initial_state=path[0], intermediate_states=states, final_state=states[-1],
+        prompt=render_prompt(pairs, path[0], depth),
+    )
+
+
+def validate_example(example: PointerExample) -> None:
+    """Reparse the actual prompt and independently recompute every saved target."""
+    if example.schema_version not in (1, 2) or example.family != "pointer":
+        raise ValueError("Unsupported pointer record schema/family")
+    if type(example.task_depth) is not int or not 1 <= example.task_depth <= (25 if example.schema_version == 1 else 256):
+        raise ValueError("Invalid task depth")
+    if any(len(pair) != 2 for pair in example.mapping):
+        raise ValueError("Every mapping entry must have two symbols")
+    mapping = dict(example.mapping)
+    if len(example.mapping) != len(SYMBOLS) or set(mapping) != set(SYMBOLS):
+        raise ValueError("Each vocabulary symbol must have exactly one rule")
+    if not set(mapping.values()) <= set(SYMBOLS):
+        raise ValueError("Mapping target outside vocabulary")
+    if example.mapping_sha256 != mapping_fingerprint(mapping):
+        raise ValueError("Mapping fingerprint mismatch")
+    if example.prompt != render_prompt(example.mapping, example.initial_state, example.task_depth):
+        raise ValueError("Prompt disagrees with structured task")
+    parsed = parse_mapping(example.prompt.splitlines()[0].removeprefix("Rules: "))
+    states = execute(parsed, example.initial_state, example.task_depth)
+    if states != example.intermediate_states or example.final_state != states[-1]:
+        raise ValueError("Intermediate/final labels disagree with reference execution")
+    if example.schema_version == 1 and len(set([example.initial_state, *states])) != example.task_depth + 1:
+        raise ValueError("Nominal path must not repeat a state")
+
+
+def generate_unconditioned_example(seed: int, depth: int, split: str, index: int,
+                                   graph_mode: str = "mixture") -> PointerExample:
+    """Sample the graph/start/order before using depth; permit cyclic execution.
+
+    Equal-probability mixture: random function, random permutation, full cycle.
+    All three have 26 rules; graph selection never consults requested horizon.
+    """
+    if type(depth) is not int or not 1 <= depth <= 256 or type(seed) is not int or seed < 0:
+        raise ValueError("Need nonnegative seed and depth in 1..256")
+    rng = random.Random(seed)
+    kinds = ("random_function", "permutation", "full_cycle")
+    kind = rng.choice(kinds) if graph_mode == "mixture" else graph_mode
+    if kind not in kinds:
+        raise ValueError("Unknown depth-independent graph mode")
+    if kind == "random_function":
+        mapping = {s: rng.choice(SYMBOLS) for s in SYMBOLS}
+    else:
+        order = rng.sample(SYMBOLS, len(SYMBOLS))
+        mapping = (dict(zip(SYMBOLS, order)) if kind == "permutation" else
+                   {s: order[(i + 1) % len(order)] for i, s in enumerate(order)})
+    start = rng.choice(SYMBOLS)
+    pairs = [[s, mapping[s]] for s in SYMBOLS]
+    rng.shuffle(pairs)
+    states = execute(mapping, start, depth)
+    return PointerExample(2, f"{split}-{index:06d}-{seed:016x}", "pointer", split, seed,
+        depth, pairs, mapping_fingerprint(mapping), start, states, states[-1],
+        render_prompt(pairs, start, depth))
+
+
+def orbit_structure(mapping: dict[str, str], start: str) -> dict[str, int]:
+    """Length before the first cycle and cycle period, independent of horizon."""
+    seen = {}
+    state = start
+    while state not in seen:
+        seen[state] = len(seen)
+        state = mapping[state]
+    return {"transient_length": seen[state], "cycle_period": len(seen) - seen[state]}

scripts/dataset/symbols.py

@@ -0,0 +1,48 @@
+"""Validate the actual tokenizer in the exact pointer prompt/answer contexts."""
+
+from typing import Any
+import re
+
+from scripts.dataset.pointer import SYMBOLS
+
+
+MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
+MODEL_REVISION = "7ae557604adf67be50417f59c2c2f167def9a775"
+
+
+def validate_prompt_tokens(tokenizer: Any, prompt: str, token_ids: dict[str, int]) -> int:
+    """Check every space-prefixed symbol span; return the answer readout index.
+
+    Prompts are raw text, with no chat template or added special tokens. The
+    answer position is the final ':' token's index (a next-token readout).
+    """
+    if not prompt.endswith("\nAnswer:"):
+        raise ValueError("Prompt must end with '\\nAnswer:' without a trailing space")
+    encoded = tokenizer(prompt, add_special_tokens=False, return_offsets_mapping=True)
+    spans = dict(zip(map(tuple, encoded["offset_mapping"]), encoded["input_ids"]))
+    for match in re.finditer(r" [A-Z](?=[,)\n])", prompt):
+        if spans.get(match.span()) != token_ids[match.group()[1:]]:
+            raise ValueError(f"Symbol does not retain its token at {match.span()}: {match.group()!r}")
+    return len(encoded["input_ids"]) - 1
+
+
+def validate_symbols(tokenizer: Any) -> dict[str, int]:
+    """Require one unique token per ' '+symbol in rules, starts, and answers."""
+    token_ids = {}
+    for symbol in SYMBOLS:
+        ids = tokenizer.encode(" " + symbol, add_special_tokens=False)
+        if len(ids) != 1 or tokenizer.decode(ids) != " " + symbol:
+            raise ValueError(f"Symbol {symbol!r} is not one reversible space-prefixed token")
+        token_ids[symbol] = ids[0]
+    if len(set(token_ids.values())) != len(SYMBOLS):
+        raise ValueError("Answer tokens must be unique")
+    # Exhaustive pair contexts, and all 26 possible answer continuations.
+    for source in SYMBOLS:
+        prompt = "Rules: " + " ".join(f"( {source}, {target})" for target in SYMBOLS)
+        prompt += f"\nStart: {source}\nSteps: 1\nAnswer:"
+        validate_prompt_tokens(tokenizer, prompt, token_ids)
+        prefix = tokenizer.encode(prompt, add_special_tokens=False)
+        for target in SYMBOLS:
+            if tokenizer.encode(prompt + " " + target, add_special_tokens=False) != [*prefix, token_ids[target]]:
+                raise ValueError(f"Answer boundary merges for {target!r}")
+    return token_ids

scripts/dataset/terminal.py

@@ -0,0 +1,31 @@
+"""Deterministic, evaluation-only absorbing-terminal variants of pointer tasks."""
+
+from dataclasses import replace
+
+from .pointer import PointerExample, execute, mapping_fingerprint, render_prompt, validate_example
+
+
+TRANSFORM_VERSION = "absorbing-terminal-v1"
+
+
+def absorbing_terminal(example: PointerExample) -> PointerExample:
+    """Replace only the final state's outgoing edge, preserving nominal execution.
+
+    The original d+1 path states are distinct, so that edge is not followed until
+    loop d+1. The changed prompt is a new distribution, not a retention label
+    retroactively applied to the original task. Source IDs/seeds support pairing.
+    """
+    validate_example(example)
+    if example.final_state in [example.initial_state, *example.intermediate_states[:-1]]:
+        raise ValueError("Terminal transform requires a final state not visited earlier; cyclic tasks need exact-step evaluation")
+    final = example.final_state
+    pairs = [[source, final if source == final else target] for source, target in example.mapping]
+    result = replace(
+        example, mapping=pairs, mapping_sha256=mapping_fingerprint(dict(pairs)),
+        prompt=render_prompt(pairs, example.initial_state, example.task_depth),
+    )
+    validate_example(result)
+    states = execute(dict(pairs), result.initial_state, result.task_depth + 2)
+    if states != [*example.intermediate_states, final, final]:
+        raise ValueError("Terminal transform changed the nominal trajectory")
+    return result

scripts/demo/__init__.py

@@ -0,0 +1 @@
+"""Public pointer demo: validated inputs and observations of native inference."""

scripts/demo/__main__.py

@@ -0,0 +1,46 @@
+"""Download/run the final checkpoint, printing one JSON snapshot per readout."""
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+import torch
+from huggingface_hub import snapshot_download
+
+from scripts.demo.inference import DemoRequest, stream_inference
+from scripts.recurrent_qwen.checkpoint import load_recurrent_checkpoint
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    source = parser.add_mutually_exclusive_group(required=True)
+    source.add_argument("--checkpoint", type=Path)
+    source.add_argument("--model", help="Hugging Face model repository ID")
+    parser.add_argument("--revision", default="main")
+    parser.add_argument("--request", required=True, type=Path, help="JSON exported from the demo")
+    parser.add_argument("--device", choices=("cpu", "cuda", "mps"), default="cpu")
+    parser.add_argument("--download", action="store_true", help="Allow explicit Hub/base downloads")
+    parser.add_argument("--max-steps", type=int, default=256)
+    parser.add_argument("--loop-cap", type=int, default=272)
+    args = parser.parse_args()
+    request = DemoRequest.parse(json.loads(args.request.read_text()), max_steps=args.max_steps)
+    if args.loop_cap < request.steps:
+        parser.error("--loop-cap must cover the requested steps")
+    path = args.checkpoint or Path(snapshot_download(args.model, revision=args.revision,
+                                                   local_files_only=not args.download))
+    torch.set_num_threads(4)
+    torch.use_deterministic_algorithms(True)
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    model, tokenizer, spec = load_recurrent_checkpoint(path, device=args.device, download=args.download)
+    model.config._attn_implementation = "sdpa"
+    model.requires_grad_(False)
+    for frame in stream_inference(model, tokenizer, spec, request,
+                                   checkpoint_id=args.model or str(path), loop_cap=args.loop_cap):
+        print(json.dumps(frame, allow_nan=False), flush=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/demo/gradio_app.py

@@ -0,0 +1,21 @@
+"""The same small queued streaming API for Hugging Face and local serving."""
+
+from typing import Callable
+
+import gradio as gr
+
+
+def create_app(predict: Callable, *, model_url: str | None = None) -> gr.Blocks:
+    with gr.Blocks(title="LoopFormer · live inference", analytics_enabled=False) as app:
+        gr.Markdown("# LoopFormer\nFollow a graph one recurrent loop at a time. "
+                    "This endpoint streams real model predictions; it does not substitute reference answers.")
+        links = "[Open the interactive playground](https://mukundshankar-dev.github.io/projects/loopformer/demo/)"
+        if model_url:
+            links += f" · [Download the model]({model_url})"
+        gr.Markdown(links)
+        request_input = gr.JSON(label="Graph request: mapping, start, steps")
+        result = gr.JSON(label="Streamed trace")
+        run = gr.Button("Run model", variant="primary")
+        run.click(predict, inputs=request_input, outputs=result, api_name="predict",
+                  concurrency_limit=1, stream_every=0.05)
+    return app.queue(max_size=16, default_concurrency_limit=1)

scripts/demo/inference.py

@@ -0,0 +1,151 @@
+"""Stream native loop observations without reference-state feedback."""
+
+from dataclasses import dataclass
+import math
+from queue import Queue
+from threading import Event, Thread
+from time import perf_counter
+from typing import Any, Iterator
+
+import torch
+
+from scripts.dataset.pointer import SYMBOLS, execute, render_prompt
+
+
+SAFETY_CAP = 272
+STOP_THRESHOLD = 0.5
+
+
+@dataclass(frozen=True)
+class DemoRequest:
+    """Ordered full A–Z table; order is part of the model's actual prompt."""
+
+    pairs: tuple[tuple[str, str], ...]
+    start: str
+    steps: int
+
+    @classmethod
+    def parse(cls, value: Any, *, max_steps: int = 256) -> "DemoRequest":
+        if not isinstance(value, dict) or set(value) != {"mapping", "start", "steps"}:
+            raise ValueError("Supply exactly mapping, start, and steps")
+        pairs = value["mapping"]
+        if not isinstance(pairs, list) or len(pairs) != 26:
+            raise ValueError("Supply all 26 A–Z rules, with one destination per node")
+        if any(not isinstance(p, list) or len(p) != 2 or
+               any(not isinstance(s, str) or s not in SYMBOLS for s in p) for p in pairs):
+            raise ValueError("Each rule must be a pair of uppercase A–Z letters")
+        if {p[0] for p in pairs} != set(SYMBOLS):
+            raise ValueError("Each A–Z source must occur exactly once")
+        if not isinstance(value["start"], str) or value["start"] not in SYMBOLS:
+            raise ValueError("Start must be an uppercase A–Z letter")
+        if type(value["steps"]) is not int or not 1 <= value["steps"] <= max_steps:
+            raise ValueError(f"Steps must be an integer from 1 to {max_steps}")
+        return cls(tuple(tuple(p) for p in pairs), value["start"], value["steps"])
+
+    def payload(self) -> dict:
+        return {"mapping": [list(p) for p in self.pairs], "start": self.start, "steps": self.steps}
+
+    def prompt(self) -> str:
+        return render_prompt([list(p) for p in self.pairs], self.start, self.steps)
+
+
+class InferenceCancelled(Exception):
+    """The client closed its generator; stop at the next completed readout."""
+
+
+def score_trace(request: DemoRequest, rows: list[dict], stop_reason: str) -> dict:
+    """Final letter, strict requested prefix and timing are separate metrics."""
+    expected = execute(dict(request.pairs), request.start, request.steps)
+    predictions = [row["prediction"] for row in rows]
+    exact_stop = stop_reason == "head" and len(rows) == request.steps
+    correct = bool(predictions) and predictions[-1] == expected[-1]
+    complete = len(predictions) >= request.steps and predictions[:request.steps] == expected
+    return {"answer_correct": correct, "complete_trajectory": complete,
+            "exact_stop": exact_stop, "joint_success": correct and exact_stop,
+            "stop_reason": stop_reason, "executed_loops": len(rows),
+            "first_error": next((r["loop"] for r in rows[:request.steps] if not r["correct"]), None)}
+
+
+def stream_inference(model: Any, tokenizer: Any, spec: dict, request: DemoRequest,
+                     *, checkpoint_id: str, cancel: Event | None = None,
+                     loop_cap: int | None = None) -> Iterator[dict]:
+    """Yield cumulative JSON snapshots as loops arrive, then a final summary.
+
+    One worker owns one native forward. The bounded 272-loop cap also bounds
+    the queue. A consumer must serialize calls to its model (Gradio concurrency
+    one). Closing the iterator cancels and joins the worker before releasing GPU
+    ownership. Targets are computed only in the consumer, never passed to R.
+    """
+    if model.training or model.completion_head is None:
+        raise ValueError("Demo requires an eval-mode model with learned completion")
+    loop_cap = SAFETY_CAP if loop_cap is None else loop_cap
+    if type(loop_cap) is not int or loop_cap < 1:
+        raise ValueError("Loop cap must be a positive integer")
+    cancel = cancel if cancel is not None else Event()
+    queue: Queue = Queue()
+    rows: list[dict] = []
+    mapping = dict(request.pairs)
+    reference = request.start
+    previous_prediction = request.start
+    started = perf_counter()
+    metadata = {"schema_version": 1, "source": "live", "checkpoint": checkpoint_id,
+                "request": request.payload(), "threshold": STOP_THRESHOLD, "safety_cap": loop_cap}
+
+    def observe(loop, scores, stop_logits, state_logits):
+        if cancel.is_set():
+            raise InferenceCancelled()
+        scores = scores[0].float().cpu()
+        probability = float(stop_logits.sigmoid().item())
+        if not torch.isfinite(scores).all() or not math.isfinite(probability):
+            raise FloatingPointError("Nonfinite model output")
+        top = scores.topk(3)
+        probabilities = scores.softmax(-1)
+        queue.put(("loop", {"loop": loop, "prediction": SYMBOLS[int(scores.argmax())],
+                   "direct_prediction": SYMBOLS[int(state_logits[0].argmax())] if state_logits is not None else None,
+                   "stop_probability": probability,
+                   "top3": [{"symbol": SYMBOLS[i], "probability": float(probabilities[i])}
+                            for i in top.indices.tolist()],
+                   "arrival_seconds": perf_counter() - started}))
+
+    def run():
+        try:
+            if cancel.is_set():
+                raise InferenceCancelled()
+            device = next(model.parameters()).device
+            encoded = tokenizer(request.prompt(), add_special_tokens=False, return_tensors="pt")
+            with torch.inference_mode():
+                result = model(input_ids=encoded["input_ids"].to(device),
+                               attention_mask=encoded["attention_mask"].to(device),
+                               num_loops=loop_cap, completion_threshold=STOP_THRESHOLD,
+                               readout_token_ids=spec["token_ids"], loop_callback=observe)
+            signaled = result.stop_logits[0, -1].sigmoid().item() >= STOP_THRESHOLD
+            queue.put(("done", "head" if signaled else "cap"))
+        except InferenceCancelled:
+            queue.put(("cancelled", None))
+        except Exception as error:
+            # Transport the original failure to the owning generator, never hide it.
+            queue.put(("error", error))
+
+    worker = Thread(target=run, name="loopformer-inference", daemon=True)
+    worker.start()
+    try:
+        while True:
+            kind, value = queue.get()
+            if kind == "error":
+                raise value
+            if kind == "cancelled":
+                return
+            if kind == "done":
+                yield {**metadata, "status": "complete", "trace": list(rows),
+                       "summary": score_trace(request, rows, value),
+                       "inference_seconds": perf_counter() - started}
+                return
+            reference = mapping[reference]
+            value.update(expected=reference, correct=value["prediction"] == reference,
+                         edge_valid=mapping[previous_prediction] == value["prediction"])
+            previous_prediction = value["prediction"]
+            rows.append(value)
+            yield {**metadata, "status": "streaming", "trace": list(rows)}
+    finally:
+        cancel.set()
+        worker.join()

scripts/recurrent_qwen/__init__.py

@@ -0,0 +1,6 @@
+"""Shared-depth Qwen model mechanics, without import-time loading."""
+
+from .model import RecurrentQwen
+from .outputs import RecurrentOutput
+
+__all__ = ["RecurrentQwen", "RecurrentOutput"]

scripts/recurrent_qwen/checkpoint.py

@@ -0,0 +1,124 @@
+"""Portable trainable recurrent tensors plus explicit architecture/base identity."""
+
+import json
+from pathlib import Path
+from typing import Any
+
+import torch
+from transformers import AutoTokenizer, Qwen2ForCausalLM
+
+from .lora_utils import attach_recurrent_lora
+from .model import RecurrentQwen
+from scripts.dataset.pointer import SYMBOLS
+
+
+EXECUTOR_FORMAT = "loopformer-executor-v2"
+FORMAT = "loopformer-stage1-v1"
+COMPLETION_FORMAT = "loopformer-stage1-completion-v1"
+FIXED_PROMPT_FORMAT = "loopformer-stage1-fixed-prompt-v1"
+
+
+def checkpoint_mode(spec: dict) -> str:
+    """Legacy checkpoints mean full-sequence recurrence; never infer new modes."""
+    mode = spec.get("recurrence_mode", "full_sequence")
+    if spec.get("format") not in (FORMAT, COMPLETION_FORMAT, FIXED_PROMPT_FORMAT, EXECUTOR_FORMAT):
+        raise ValueError("Unsupported recurrent checkpoint format")
+    if (spec["format"] == EXECUTOR_FORMAT) != ("executor" in spec):
+        raise ValueError("Executor architecture requires the versioned executor checkpoint format")
+    expected = "fixed_prompt" if spec["format"] in (FIXED_PROMPT_FORMAT, EXECUTOR_FORMAT) else "full_sequence"
+    if mode != expected:
+        raise ValueError("Checkpoint format and recurrence_mode disagree")
+    return mode
+
+
+def cpu_tree(value: Any) -> Any:
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu().clone()
+    if isinstance(value, dict):
+        return {key: cpu_tree(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [cpu_tree(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(cpu_tree(item) for item in value)
+    return value
+
+
+def save_checkpoint(path: Path, model: RecurrentQwen, tokenizer: Any, spec: dict, training_state: dict) -> None:
+    """Write a new checkpoint directory; metadata is the final completion marker."""
+    if spec.get("recurrence_mode", "full_sequence") != model.recurrence_mode:
+        raise ValueError("Checkpoint spec and recurrence_mode disagree")
+    weights = {name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad}
+    completion = model.completion_head is not None
+    expected_completion = spec.get("completion_head")
+    if completion != (expected_completion is not None):
+        raise ValueError("Checkpoint spec and completion head disagree")
+    if completion and expected_completion != {"intermediate": model.completion_head.intermediate}:
+        raise ValueError("Checkpoint completion-head architecture differs from spec")
+    new_architecture = "executor" in spec
+    if not new_architecture and (not weights or any(not ((name.startswith("recurrent.") and ".lora_" in name) or
+                               (completion and name.startswith("completion_head."))) for name in weights)):
+        raise ValueError("Checkpoint expects recurrent LoRA and optional completion head only")
+    if new_architecture:
+        if spec["executor"]["train_scope"] != model.train_scope:
+            raise ValueError("Checkpoint training scope differs from model")
+        if not weights:
+            raise ValueError("Cannot save an empty executor")
+    path.mkdir(parents=True, exist_ok=False)
+    torch.save(cpu_tree(weights), path / "adapter_model.pt")
+    torch.save(cpu_tree(training_state), path / "training_state.pt")
+    tokenizer.save_pretrained(path)
+    version = FIXED_PROMPT_FORMAT if model.recurrence_mode == "fixed_prompt" else (COMPLETION_FORMAT if completion else FORMAT)
+    if new_architecture:
+        version = EXECUTOR_FORMAT
+    (path / "recurrent_config.json").write_text(json.dumps({**spec, "format": version}, indent=2) + "\n")
+
+
+def restore_adapters(model: RecurrentQwen, path: Path) -> None:
+    if checkpoint_mode(json.loads((path / "recurrent_config.json").read_text())) != model.recurrence_mode:
+        raise ValueError("Adapter checkpoint recurrence_mode differs from model")
+    state = torch.load(path / "adapter_model.pt", map_location="cpu", weights_only=True)
+    expected = {name: value for name, value in model.named_parameters() if value.requires_grad}
+    if set(state) != set(expected):
+        raise ValueError("Adapter checkpoint names do not match the recurrent model")
+    with torch.no_grad():
+        for name, parameter in expected.items():
+            tensor = state[name]
+            if tensor.shape != parameter.shape or not torch.isfinite(tensor).all():
+                raise ValueError(f"Invalid adapter tensor: {name}")
+            parameter.copy_(tensor)
+
+
+def load_recurrent_checkpoint(path: Path, *, device: str = "cpu", dtype: str = "float32",
+                              download: bool = False) -> tuple[RecurrentQwen, Any, dict]:
+    spec = json.loads((path / "recurrent_config.json").read_text())
+    mode = checkpoint_mode(spec)
+    if spec.get("prompt_format") != "dataset_raw" or spec.get("loss_vocabulary") != "symbols":
+        raise ValueError("Unsupported recurrent checkpoint format/prompt/loss")
+    completion = spec["format"] == COMPLETION_FORMAT or (mode == "fixed_prompt" and "completion_head" in spec)
+    if completion and (not isinstance(spec.get("completion_head"), dict) or
+                       type(spec["completion_head"].get("intermediate")) is not int or
+                       spec["completion_head"]["intermediate"] < 1):
+        raise ValueError("Invalid saved completion-head architecture")
+    if not completion and "completion_head" in spec:
+        raise ValueError("Legacy checkpoint cannot contain a completion head")
+    if spec.get("symbols") != list(SYMBOLS) or len(set(spec.get("token_ids", []))) != len(SYMBOLS):
+        raise ValueError("Checkpoint requires the complete A–Z symbol vocabulary")
+    tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=True, trust_remote_code=False)
+    for symbol, token in zip(spec["symbols"], spec["token_ids"], strict=True):
+        if tokenizer.encode(" " + symbol, add_special_tokens=False) != [token]:
+            raise ValueError("Checkpoint tokenizer and symbol IDs disagree")
+    base = Qwen2ForCausalLM.from_pretrained(
+        spec["base_model"], revision=spec["revision"], local_files_only=not download,
+        dtype=getattr(torch, dtype), attn_implementation="eager", trust_remote_code=False,
+    ).to(device).eval()
+    model = RecurrentQwen(base, spec["recurrent_start"], spec["recurrent_end"], recurrence_mode=mode)
+    executor = spec.get("executor")
+    if executor is None or executor["train_scope"] == "lora":
+        attach_recurrent_lora(model, rank=spec["lora_rank"], alpha=spec["lora_alpha"])
+    if executor is not None:
+        model.configure_executor(tokenizer, **executor)
+    if completion and model.completion_head is None:
+        model.enable_completion(spec["completion_head"]["intermediate"])
+    restore_adapters(model, path)
+    model.eval()
+    return model, tokenizer, spec

scripts/recurrent_qwen/completion.py

@@ -0,0 +1,26 @@
+"""Hidden-state-only completion head for explicit-step pointer recurrence."""
+
+import torch
+from torch import Tensor, nn
+
+
+class CompletionHead(nn.Module):
+    """One stop logit per answer-position recurrent state [B,H].
+
+    No numeric loop index, parsed requested depth, decoded symbol, or target is
+    accepted by this interface. The prompt can encode its requested Steps.
+    """
+
+    def __init__(self, width: int, intermediate: int = 128) -> None:
+        super().__init__()
+        if width < 1 or intermediate < 1:
+            raise ValueError("Completion-head widths must be positive")
+        self.width = width
+        self.intermediate = intermediate
+        self.network = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, intermediate),
+                                     nn.GELU(), nn.Linear(intermediate, 1))
+
+    def forward(self, answer_state: Tensor) -> Tensor:
+        if answer_state.ndim != 2 or answer_state.shape[-1] != self.width:
+            raise ValueError("Completion head requires answer states [batch, hidden]")
+        return self.network(answer_state).squeeze(-1)

scripts/recurrent_qwen/interfaces.py

@@ -0,0 +1,139 @@
+"""Explicit text routing and learned interfaces for isolated pointer execution."""
+from functools import lru_cache
+import re
+from typing import Any
+
+import torch
+from torch import Tensor, nn
+
+from .number_reader import SharedNumberReader
+
+
+class PromptRouter:
+    """Remove the Steps line without parsing its value or changing public inputs.
+
+    Preserve the original answer index for existing diagnostic callers. Padding
+    inserted before that position is masked and RoPE positions are compacted, so
+    neither count digits nor their token length reach the executor. Cached CPU
+    tokenization is bounded; no prompt or activation survives in a checkpoint.
+    """
+    def __init__(self, tokenizer: Any) -> None:
+        self.tokenizer = tokenizer
+        self.executor_ids = lru_cache(maxsize=4096)(self._executor_ids)
+
+    def _executor_ids(self, ids: tuple[int, ...]) -> tuple[int, ...]:
+        prompt = self.tokenizer.decode(list(ids), skip_special_tokens=False,
+                                       clean_up_tokenization_spaces=False)
+        if not re.fullmatch(r"Rules: [^\n]+\nStart: [A-Z]\nSteps: [1-9][0-9]*\nAnswer:", prompt):
+            raise ValueError("Isolated execution requires the raw Rules/Start/Steps/Answer prompt")
+        routed = re.sub(r"\nSteps: [1-9][0-9]*", "", prompt)
+        return tuple(self.tokenizer.encode(routed, add_special_tokens=False))
+
+    def __call__(self, ids: Tensor, mask: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        rows = ids.detach().cpu().tolist()
+        masks = mask.detach().cpu().tolist()
+        routed = torch.full_like(ids, self.tokenizer.pad_token_id)
+        attention = torch.zeros_like(mask)
+        for row, (tokens, valid) in enumerate(zip(rows, masks, strict=True)):
+            positions = [i for i, active in enumerate(valid) if active]
+            executor = self.executor_ids(tuple(tokens[i] for i in positions))
+            end = positions[-1]
+            if not executor or len(executor) > end + 1:
+                raise ValueError("Routed executor prompt does not fit original input positions")
+            routed[row, :len(executor) - 1] = torch.tensor(executor[:-1], device=ids.device)
+            routed[row, end] = executor[-1]
+            attention[row, :len(executor) - 1] = 1
+            attention[row, end] = 1
+        positions = (attention.cumsum(-1) - 1).clamp_min(0)
+        return routed, attention, positions
+
+
+class ReentryBridge(nn.Module):
+    """Learned map from R output back to R input, applied only on re-entry.
+
+    Normalize direction, then restore the *initial* working-state RMS scale.
+    The scale contains rules/start information only and stays differentiable.
+    Identity initialization of the projection does not claim T>1 equivalence.
+    """
+    def __init__(self, width: int) -> None:
+        super().__init__()
+        self.projection = nn.Linear(width, width, bias=False)
+        nn.init.eye_(self.projection.weight)
+        self.gain = nn.Parameter(torch.ones(width))
+
+    def forward(self, state: Tensor, initial: Tensor) -> Tensor:
+        rms = state.float().square().mean(-1, keepdim=True).add(1e-6).sqrt()
+        scale = initial.float().square().mean(-1, keepdim=True).add(1e-6).sqrt()
+        normalized = (state.float() / rms * scale).to(state.dtype)
+        return self.projection(normalized * self.gain)
+
+
+class RecurrentController(nn.Module):
+    """Own learned memory; detached prompt/executor features, no numeric clock."""
+    kind = 'gru'
+    def __init__(self, width: int, intermediate: int = 128) -> None:
+        super().__init__()
+        self.intermediate = intermediate
+        self.context = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, intermediate), nn.Tanh())
+        self.observation = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, intermediate), nn.Tanh())
+        self.cell = nn.GRUCell(intermediate, intermediate)
+        self.readout = nn.Linear(intermediate, 1)
+
+    def initialize(self, prompt_state: Tensor) -> Tensor:
+        return self.context(prompt_state.detach())
+
+    def advance(self, state: Tensor, memory: Tensor) -> tuple[Tensor, Tensor]:
+        memory = self.cell(self.observation(state.detach()), memory)
+        return self.readout(memory).squeeze(-1), memory
+
+
+class AffineSuffixController(RecurrentController):
+    """Learned prompt reading and unbounded scalar recurrence; no parsed clock.
+
+    Initialization consumes flattened final-eight frozen token embeddings.
+    The trainable cell implements a*m+b, initialized to identity (b=0), never
+    a programmed decrement. Numerical labels train the cell through free-running
+    recurrence; stop BCE trains its readout without distorting counting dynamics.
+    Executor observations are intentionally ignored for this requested-count
+    controller: timing is independent of the pointer's identity or correctness.
+    """
+    kind = 'affine_suffix'
+    suffix_tokens = 8
+
+    def __init__(self, width: int) -> None:
+        nn.Module.__init__(self)
+        self.width = width
+        self.intermediate = 1
+        self.context = nn.Linear(width * self.suffix_tokens, 1)
+        self.observation = nn.Identity()
+        self.cell = nn.Linear(1, 1)
+        self.readout = nn.Linear(1, 1)
+        nn.init.ones_(self.cell.weight)
+        nn.init.zeros_(self.cell.bias)
+        nn.init.zeros_(self.readout.weight)
+        nn.init.zeros_(self.readout.bias)
+
+    def advance(self, state: Tensor, memory: Tensor) -> tuple[Tensor, Tensor]:
+        self.observation(state.detach())  # Preserve observation hooks for cache/native checks.
+        memory = self.cell(memory)
+        return self.readout(memory.detach()).squeeze(-1), memory
+
+
+class SharedNumberController(AffineSuffixController):
+    """Same learned countdown; share the number reader across positions/lengths."""
+    kind = 'shared_number'
+
+    def __init__(self, width: int) -> None:
+        super().__init__(width)
+        self.context = SharedNumberReader(width)
+
+
+def controller_head(width: int, intermediate: int, kind: str = 'gru') -> RecurrentController:
+    """Explicit architecture factory; historical checkpoints default to GRU."""
+    if kind == 'gru':
+        return RecurrentController(width, intermediate)
+    if kind == 'affine_suffix' and intermediate == 1:
+        return AffineSuffixController(width)
+    if kind == 'shared_number' and intermediate == 1:
+        return SharedNumberController(width)
+    raise ValueError('Unknown controller kind or incompatible memory width')

scripts/recurrent_qwen/lora_utils.py

@@ -0,0 +1,35 @@
+"""Attach PEFT LoRA only to the shared recurrent decoder group."""
+
+from peft import LoraConfig, inject_adapter_in_model
+
+from .model import RecurrentQwen
+
+
+def attach_recurrent_lora(
+    model: RecurrentQwen,
+    *,
+    rank: int = 8,
+    alpha: int = 16,
+    targets: tuple[str, ...] = ("q_proj", "v_proj"),
+) -> None:
+    """Mutate R in place; zero-initialized B preserves the original forward.
+
+    No bias training, dropout, bridge, or adapters outside R. Original weights
+    remain frozen. An initial zero gradient for A is expected when B is zero.
+    """
+    if rank < 1 or alpha < 1:
+        raise ValueError("LoRA rank and alpha must be positive")
+    valid = {"q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"}
+    if not targets or len(set(targets)) != len(targets) or not set(targets) <= valid:
+        raise ValueError(f"Choose unique LoRA targets from {sorted(valid)}")
+    if getattr(model.recurrent, "peft_config", None):
+        raise ValueError("Recurrent LoRA is already attached")
+    config = LoraConfig(
+        r=rank,
+        lora_alpha=alpha,
+        target_modules=list(targets),
+        lora_dropout=0.0,
+        bias="none",
+        init_lora_weights=True,
+    )
+    inject_adapter_in_model(config, model.recurrent)

scripts/recurrent_qwen/memory.py

@@ -0,0 +1,35 @@
+"""Per-forward, differentiable prompt memory for a single writable position."""
+
+from dataclasses import dataclass, field
+
+import torch
+from torch import Tensor
+
+
+@dataclass
+class PromptMemory:
+    """Capture layer inputs on pass one and reuse their non-working positions.
+
+    write_mask is [B,S,1]; captured tensors are [B,S,H]. Never detach: later
+    losses must reach adapters that formed the prompt memory on the first pass.
+    Objects live for one model forward only, not across examples or updates.
+    This is an activation-memory reference path, not an attention KV cache.
+    """
+
+    write_mask: Tensor
+    layers: list[Tensor] = field(default_factory=list)
+    output: Tensor | None = None
+    keys_values: list[tuple[Tensor, Tensor]] = field(default_factory=list)
+    layer_outputs: list[Tensor] = field(default_factory=list)
+
+    def layer_input(self, index: int, hidden: Tensor) -> Tensor:
+        if self.output is None:
+            self.layers.append(hidden)
+            return hidden
+        return torch.where(self.write_mask, hidden, self.layers[index])
+
+    def block_output(self, hidden: Tensor) -> Tensor:
+        if self.output is None:
+            self.output = hidden
+            return hidden
+        return torch.where(self.write_mask, hidden, self.output)

scripts/recurrent_qwen/model.py

@@ -0,0 +1,374 @@
+"""Qwen2 decoder surgery for the pinned Transformers implementation."""
+
+import math
+from typing import Any, Callable, Literal
+
+import torch
+from torch import Tensor, nn
+from transformers import Qwen2ForCausalLM
+from transformers.masking_utils import create_causal_mask
+
+from .outputs import RecurrentOutput, answer_margin
+from .completion import CompletionHead
+from .memory import PromptMemory
+from .prefix import fixed_prefix_layer
+from torch.utils.checkpoint import checkpoint
+from .interfaces import PromptRouter, ReentryBridge, RecurrentController, AffineSuffixController, SharedNumberController, controller_head
+from .number_reader import MAX_NUMBER_TOKENS, number_token_ids
+
+
+def _answer_readout(hidden: Tensor, positions: list[int]) -> Tensor:
+    # Integer slices avoid MPS's nondeterministic scatter backward for advanced
+    # indexing/gather. Batches are deliberately tiny; positions sync once per forward.
+    return torch.stack([hidden[row, position] for row, position in enumerate(positions)])
+
+
+class DecoderBlock(nn.Module):
+    """An ordered group of original decoder objects, with no copies."""
+
+    def __init__(self, layers: list[nn.Module]) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList(layers)
+        self.prefix_reuse = False
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden: Tensor,
+        attention_mask: Tensor | None,
+        position_ids: Tensor,
+        position_embeddings: tuple[Tensor, Tensor],
+        memory: PromptMemory | None = None,
+    ) -> Tensor:
+        for index, layer in enumerate(self.layers):
+            if memory is not None:
+                hidden = memory.layer_input(index, hidden)
+            if self.prefix_reuse and memory is not None:
+                hidden = fixed_prefix_layer(layer, hidden, position_embeddings, attention_mask,
+                                            memory, index, self.gradient_checkpointing)
+                continue
+            call = layer
+            if self.gradient_checkpointing and torch.is_grad_enabled():
+                def call(x, _layer=layer, **kwargs):
+                    return checkpoint(_layer, x, use_reentrant=False, **kwargs)
+            hidden = call(
+                hidden,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                position_embeddings=position_embeddings,
+                past_key_values=None,
+                use_cache=False,
+            )
+        return memory.block_output(hidden) if memory is not None else hidden
+
+
+class RecurrentQwen(nn.Module):
+    """P -> R repeated T times -> C, reusing the supplied model's weights.
+
+    Construction freezes and shares the supplied base model in place. Load a
+    fresh base for an independent reference after adapters have been attached.
+    Only full attention and default RoPE are supported in Stage 0, as used by
+    Qwen2.5-0.5B-Instruct. Device/dtype are inherited, never selected here.
+    """
+
+    def __init__(
+        self, base: Qwen2ForCausalLM, recurrent_start: int = 6, recurrent_end: int = 18,
+        *, recurrence_mode: str = "full_sequence",
+    ) -> None:
+        super().__init__()
+        if not isinstance(base, Qwen2ForCausalLM):
+            raise TypeError("Stage 0 requires Qwen2ForCausalLM")
+        if recurrence_mode not in ("full_sequence", "fixed_prompt"):
+            raise ValueError("recurrence_mode must be full_sequence or fixed_prompt")
+        self.recurrence_mode = recurrence_mode
+        layers = list(base.model.layers)
+        if not 0 <= recurrent_start < recurrent_end <= len(layers):
+            raise ValueError("Require 0 <= recurrent_start < recurrent_end <= layer count")
+        if any(kind != "full_attention" for kind in base.config.layer_types):
+            raise ValueError("Stage 0 supports full-attention Qwen layers only")
+        if base.config.rope_parameters["rope_type"] != "default":
+            raise ValueError("Stage 0 supports default RoPE only")
+        if getattr(base, "peft_config", None):
+            raise ValueError("Supply an unadapted base; attach LoRA to recurrence afterwards")
+        base.requires_grad_(False)
+        base.zero_grad(set_to_none=True)
+        base.config.use_cache = False
+        self.config = base.config
+        self.recurrent_start = recurrent_start
+        self.recurrent_end = recurrent_end
+        self.embed_tokens = base.model.embed_tokens
+        self.rotary_emb = base.model.rotary_emb
+        self.prelude = DecoderBlock(layers[:recurrent_start])
+        self.recurrent = DecoderBlock(layers[recurrent_start:recurrent_end])
+        self.coda = DecoderBlock(layers[recurrent_end:])
+        self.norm = base.model.norm
+        self.lm_head = base.lm_head
+        self.completion_head: CompletionHead | RecurrentController | None = None
+        self.router = None
+        self.bridge = None
+        self.state_head = None
+        self.train_scope = "lora"
+        self.prefix_reuse = False
+        self.gradient_checkpointing = False
+        self.train(base.training)
+
+    def enable_completion(self, intermediate: int = 128) -> None:
+        """Attach a trainable stop readout without changing recurrent inputs."""
+        if self.completion_head is not None:
+            raise ValueError("Completion head is already attached")
+        self.completion_head = CompletionHead(self.config.hidden_size, intermediate).to(
+            device=self.embed_tokens.weight.device, dtype=self.embed_tokens.weight.dtype)
+
+    def configure_executor(self, tokenizer: Any, *, train_scope: str = "recurrent",
+                           bridge: bool = True, direct_readout: bool = True,
+                           controller_size: int = 128, isolated: bool = True,
+                           prefix_reuse: bool = False, gradient_checkpointing: bool = False,
+                           controller_kind: str = 'gru') -> None:
+        """Opt-in architecture; legacy forwards/checkpoints keep their semantics."""
+        if train_scope not in ("lora", "recurrent", "full"):
+            raise ValueError("train_scope must be lora, recurrent, or full")
+        if self.recurrence_mode != "fixed_prompt":
+            raise ValueError("Executor interfaces require fixed_prompt recurrence")
+        if train_scope != "lora" and any(".lora_" in name for name, _ in self.named_parameters()):
+            raise ValueError("Full-weight scopes require an unadapted base, not an attached LoRA model")
+        if prefix_reuse and self.config.attention_dropout != 0:
+            raise ValueError("Prefix reuse requires zero attention dropout for equivalent execution")
+        self.train_scope = train_scope
+        if train_scope == "recurrent":
+            self.recurrent.requires_grad_(True)
+        elif train_scope == "full":
+            self.requires_grad_(True)
+        device, dtype = self.embed_tokens.weight.device, self.embed_tokens.weight.dtype
+        if isolated:
+            self.router = PromptRouter(tokenizer)
+            if controller_size:
+                self.completion_head = controller_head(self.config.hidden_size, controller_size, controller_kind).to(device=device, dtype=dtype)
+        if bridge:
+            self.bridge = ReentryBridge(self.config.hidden_size).to(device=device, dtype=dtype)
+        if direct_readout:
+            self.state_head = nn.Linear(self.config.hidden_size, 26).to(device=device, dtype=dtype)
+        self.prefix_reuse = prefix_reuse
+        self.gradient_checkpointing = gradient_checkpointing
+        for block in (self.prelude, self.recurrent, self.coda):
+            block.prefix_reuse = prefix_reuse
+            block.gradient_checkpointing = gradient_checkpointing
+
+    def controller_prompt_state(self, input_ids: Tensor, attention_mask: Tensor,
+                                answer_positions: Tensor | None = None) -> Tensor:
+        """Frozen native initializer features [B,H_context], without R/C execution.
+
+        This interface runs neither R nor C and consumes only token IDs/mask.
+        GRU uses full-prompt P's answer position; affine_suffix uses the final
+        eight frozen token embeddings. Shared-number uses routed field-token
+        embeddings and validity masks. No branch interprets a numeric value.
+        """
+        if input_ids.ndim != 2 or attention_mask.shape != input_ids.shape or not attention_mask.any(-1).all():
+            raise ValueError('Need aligned nonempty [B,S] prompt IDs and mask')
+        if answer_positions is None:
+            answer_positions = (attention_mask.long() * torch.arange(1, input_ids.shape[1] + 1,
+                                device=input_ids.device)).argmax(-1)
+        if isinstance(self.completion_head, SharedNumberController):
+            rows = []
+            for ids, mask in zip(input_ids.detach().cpu().tolist(), attention_mask.detach().cpu().tolist(), strict=True):
+                valid = tuple(t for t, active in zip(ids, mask, strict=True) if active)
+                rows.append(number_token_ids(self.router.tokenizer, valid))
+            with torch.no_grad():
+                context = self.embed_tokens.weight.new_zeros(len(rows), MAX_NUMBER_TOKENS, self.config.hidden_size + 1)
+                for i, tokens in enumerate(rows):
+                    context[i, :len(tokens), :-1] = self.embed_tokens(torch.tensor(tokens, device=input_ids.device))
+                    context[i, :len(tokens), -1] = 1
+            return context.flatten(1)
+        if isinstance(self.completion_head, AffineSuffixController):
+            k = self.completion_head.suffix_tokens
+            if (answer_positions < k - 1).any():
+                raise ValueError('Affine suffix initialization requires eight final prompt tokens')
+            with torch.no_grad():
+                embeddings = self.embed_tokens(input_ids)
+                return torch.stack([embeddings[i, end - k + 1:end + 1].flatten()
+                                    for i, end in enumerate(answer_positions.tolist())])
+        with torch.no_grad():
+            context = self.embed_tokens(input_ids)
+            positions = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0)
+            mask = create_causal_mask(config=self.config, inputs_embeds=context,
+                attention_mask=attention_mask, past_key_values=None, position_ids=positions)
+            context = self.prelude(context, mask, positions, self.rotary_emb(context, positions))
+        return _answer_readout(context, answer_positions.tolist())
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        attention_mask: Tensor | None = None,
+        *,
+        num_loops: int = 1,
+        position_ids: Tensor | None = None,
+        answer_positions: Tensor | None = None,
+        labels: Tensor | None = None,
+        allowed_token_ids: Tensor | None = None,
+        return_hidden_states: bool = False,
+        logits_mode: Literal["answer", "all"] = "answer",
+        stop_policy: Callable[[int, Tensor, Tensor], bool] | None = None,
+        completion_threshold: float | None = None,
+        readout_token_ids: list[int] | None = None,
+        loop_callback: Callable[[int, Tensor, Tensor | None, Tensor | None], None] | None = None,
+    ) -> RecurrentOutput:
+        """Run complete prompt states through shared recurrence, without KV cache.
+
+        Inputs/mask: [B,S]; positions: [1,S] or [B,S], default arange(S),
+        matching ordinary Qwen forward (including padding). Answer positions:
+        [B], default last unmasked token. Labels: [B,T] *token IDs*, one
+        target per loop, with an explicit allowed answer set [A]. No CE loss
+        or causal label shifting is performed. Pass repeated final labels
+        explicitly if final-answer margins, rather than step margins, are wanted.
+        Optional stop_policy(t, h_t, answer_logits_t) is inference-only, batch 1,
+        and may return True to end after loop t. It never sees labels.
+        completion_threshold enables the checkpoint's hidden-state-only head,
+        independent of task depth, under the num_loops safety cap (batch 1).
+        fixed_prompt keeps only the final unmasked position writable; returned
+        full-shaped states contain fixed first-pass prefix outputs for diagnostics.
+        loop_callback receives one-based loop, C logits [1,A], optional stop
+        logits [1], and optional R logits [1,26] after each readout, including
+        the stopping loop. Detached copies protect inference from observer edits.
+        It is inference-only; its return value cannot choose when to stop.
+        """
+        if type(num_loops) is not int or num_loops < 1:
+            raise ValueError("num_loops must be a positive integer")
+        if loop_callback is not None and (self.training or torch.is_grad_enabled() or
+                input_ids.shape[0] != 1 or labels is not None or logits_mode != "answer"):
+            raise ValueError("Loop observation requires eval/no_grad, batch 1, no labels, and answer logits")
+        if stop_policy is not None and (self.training or torch.is_grad_enabled() or input_ids.shape[0] != 1 or labels is not None or logits_mode != "answer"):
+            raise ValueError("Stopping requires eval/no_grad, batch 1, no labels, and answer logits")
+        if completion_threshold is not None:
+            if (not math.isfinite(completion_threshold) or not 0 < completion_threshold < 1 or
+                    self.completion_head is None or stop_policy is not None):
+                raise ValueError("Completion stopping requires a head, threshold in (0,1), and no other stop policy")
+            if self.training or torch.is_grad_enabled() or input_ids.shape[0] != 1 or labels is not None or logits_mode != "answer":
+                raise ValueError("Completion stopping requires eval/no_grad, batch 1, no labels, and answer logits")
+        if logits_mode not in ("answer", "all"):
+            raise ValueError("logits_mode must be 'answer' or 'all'")
+        if input_ids.ndim != 2 or input_ids.dtype != torch.long or 0 in input_ids.shape:
+            raise ValueError("input_ids must be a nonempty torch.long [batch, sequence] tensor")
+        if input_ids.device != self.embed_tokens.weight.device:
+            raise ValueError("Input IDs and model must share a device")
+        if (input_ids < 0).any() or (input_ids >= self.config.vocab_size).any():
+            raise ValueError("Input IDs must be within the model vocabulary")
+        batch, sequence = input_ids.shape
+        device = input_ids.device
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        if attention_mask.shape != input_ids.shape or attention_mask.device != device:
+            raise ValueError("attention_mask must match input_ids shape and device")
+        if not ((attention_mask == 0) | (attention_mask == 1)).all():
+            raise ValueError("attention_mask must contain only 0 and 1")
+        if not attention_mask.bool().any(dim=-1).all():
+            raise ValueError("Each example must have at least one unmasked token")
+        if position_ids is None:
+            position_ids = torch.arange(sequence, device=device).unsqueeze(0)
+        if (
+            position_ids.shape not in ((1, sequence), (batch, sequence))
+            or position_ids.dtype != torch.long
+            or position_ids.device != device
+            or (position_ids < 0).any()
+        ):
+            raise ValueError("position_ids must be nonnegative long [1,S] or [B,S] on the input device")
+        if answer_positions is None:
+            indices = torch.arange(sequence, device=device).expand(batch, -1)
+            answer_positions = indices.masked_fill(~attention_mask.bool(), -1).max(-1).values
+        if (
+            answer_positions.shape != (batch,)
+            or answer_positions.dtype != torch.long
+            or answer_positions.device != device
+            or (answer_positions < 0).any()
+            or (answer_positions >= sequence).any()
+        ):
+            raise ValueError("answer_positions must be long [batch] valid sequence indices on the input device")
+        rows = torch.arange(batch, device=device)
+        if not attention_mask[rows, answer_positions].bool().all():
+            raise ValueError("Answer positions cannot point to padding")
+        if self.recurrence_mode == "fixed_prompt":
+            last = torch.arange(sequence, device=device).expand(batch, -1).masked_fill(~attention_mask.bool(), -1).max(-1).values
+            if not torch.equal(answer_positions, last):
+                raise ValueError("fixed_prompt requires the final unmasked position as its working state")
+        readout_positions = answer_positions.tolist()
+        if (labels is None) != (allowed_token_ids is None):
+            raise ValueError("Provide both labels and allowed_token_ids to compute margins")
+        if labels is not None and labels.shape != (batch, num_loops):
+            raise ValueError("labels must have shape [batch, num_loops]")
+
+        controller_initial = None
+        if self.router is not None:
+            if logits_mode != "answer":
+                raise ValueError("Isolated execution supports answer-position logits only")
+            # Controller prompt encoding is isolated even for full-model training.
+            controller_initial = self.controller_prompt_state(input_ids, attention_mask, answer_positions)
+            input_ids, attention_mask, position_ids = self.router(input_ids, attention_mask)
+        hidden = self.embed_tokens(input_ids)
+        causal_mask = create_causal_mask(
+            config=self.config,
+            inputs_embeds=hidden,
+            attention_mask=attention_mask,
+            past_key_values=None,
+            position_ids=position_ids,
+        )
+        position_embeddings = self.rotary_emb(hidden, position_ids)
+        block_args = (causal_mask, position_ids, position_embeddings)
+        hidden = self.prelude(hidden, *block_args)
+        initial_hidden = hidden if return_hidden_states else None
+        bridge_initial = _answer_readout(hidden, readout_positions)
+        controller_memory = (self.completion_head.initialize(controller_initial)
+                             if isinstance(self.completion_head, RecurrentController) else None)
+        recurrent_memory = coda_memory = None
+        if self.recurrence_mode == "fixed_prompt":
+            write_mask = (torch.arange(sequence, device=device)[None, :] == answer_positions[:, None]).unsqueeze(-1)
+            recurrent_memory, coda_memory = PromptMemory(write_mask), PromptMemory(write_mask)
+        states, logits, margins, stop_logits, state_logits = [], [], [], [], []
+        for loop in range(num_loops):
+            if loop and self.bridge is not None:
+                working = self.bridge(_answer_readout(hidden, readout_positions), bridge_initial)
+                hidden = torch.where(write_mask, working[:, None, :], hidden)
+            hidden = (self.recurrent(hidden, *block_args, memory=recurrent_memory)
+                      if recurrent_memory is not None else self.recurrent(hidden, *block_args))
+            working = _answer_readout(hidden, readout_positions)
+            if isinstance(self.completion_head, RecurrentController):
+                stop, controller_memory = self.completion_head.advance(working, controller_memory)
+                stop_logits.append(stop)
+            elif self.completion_head is not None:
+                stop_logits.append(self.completion_head(working))
+            if self.state_head is not None:
+                state_logits.append(self.state_head(working))
+            if return_hidden_states:
+                states.append(hidden)
+            # C stays differentiable even when frozen: readout loss must reach R.
+            decoded = self.norm(self.coda(hidden, *block_args, memory=coda_memory)
+                                if coda_memory is not None else self.coda(hidden, *block_args))
+            readout = decoded if logits_mode == "all" else _answer_readout(decoded, readout_positions)
+            if readout_token_ids is None:
+                scores = self.lm_head(readout)
+            else:
+                if labels is not None:
+                    raise ValueError("Use full-vocabulary readout when requesting token-ID margins")
+                ids = torch.tensor(readout_token_ids, device=device)
+                scores = nn.functional.linear(readout, self.lm_head.weight.index_select(0, ids))
+            logits.append(scores)
+            if loop_callback is not None:
+                loop_callback(loop + 1, scores.detach().clone(),
+                              stop_logits[-1].detach().clone() if stop_logits else None,
+                              state_logits[-1].detach().clone() if state_logits else None)
+            if labels is not None:
+                answer_scores = _answer_readout(scores, readout_positions) if logits_mode == "all" else scores
+                margins.append(answer_margin(answer_scores, labels[:, loop], allowed_token_ids))
+            if stop_policy is not None and stop_policy(loop + 1, hidden, scores):
+                break
+            if completion_threshold is not None:
+                probability = stop_logits[-1].sigmoid().item()
+                if not math.isfinite(probability):
+                    raise FloatingPointError("Nonfinite completion probability")
+                if probability >= completion_threshold:
+                    break
+        return RecurrentOutput(
+            loop_logits=tuple(logits),
+            state_logits=torch.stack(state_logits, dim=1) if state_logits else None,
+            hidden_states=tuple(states) if return_hidden_states else None,
+            initial_hidden_state=initial_hidden,
+            margins=torch.stack(margins, dim=1) if margins else None,
+            stop_logits=torch.stack(stop_logits, dim=1) if stop_logits else None,
+        )

scripts/recurrent_qwen/number_reader.py

@@ -0,0 +1,61 @@
+"""Shared learned token composition, with no numeric parser in inference."""
+import re
+
+import torch
+from torch import Tensor, nn
+
+
+MAX_NUMBER_TOKENS = 8
+
+
+def number_token_ids(tokenizer, ids: tuple[int, ...]) -> tuple[int, ...]:
+    """Route only the raw Steps text; never convert it to a number.
+
+    One token per character is a tokenizer contract, not a digit-value lookup.
+    The public prompt and executor routing remain unchanged.
+    """
+    text = tokenizer.decode(list(ids), skip_special_tokens=False,
+                            clean_up_tokenization_spaces=False)
+    match = re.fullmatch(r'Rules: [^\n]+\nStart: [A-Z]\nSteps: ([1-9][0-9]*)\nAnswer:', text)
+    if match is None:
+        raise ValueError('Number reader requires the raw Rules/Start/Steps/Answer prompt')
+    field = match.group(1)
+    tokens = tokenizer.encode(field, add_special_tokens=False)
+    if not 1 <= len(tokens) <= MAX_NUMBER_TOKENS or len(tokens) != len(field):
+        raise ValueError('Number reader requires one token per decimal character and at most eight characters')
+    # Length alone does not rule out a tokenizer that groups/reorders characters.
+    if any(tokenizer.decode([t], clean_up_tokenization_spaces=False) != c
+           for t, c in zip(tokens, field, strict=True)):
+        raise ValueError('Number reader tokenizer does not preserve individual decimal characters')
+    return tuple(tokens)
+
+
+class SharedNumberReader(nn.Module):
+    """Read padded frozen embeddings [B,8*(H+1)] using the same map at each slot.
+
+    Last channel of each slot is a validity mask. Valid tokens are left aligned.
+    The gain starts at identity, not a decimal radix. Training labels fit both
+    gain and token projection. Accumulate in float64 to reduce numerical error,
+    then return the controller's scalar-memory dtype [B,1].
+    """
+    def __init__(self, width: int) -> None:
+        super().__init__()
+        self.width = width
+        self.in_features = MAX_NUMBER_TOKENS * (width + 1)
+        self.token_value = nn.Linear(width, 1)
+        self.gain = nn.Parameter(torch.ones(()))
+
+    def unpack(self, context: Tensor) -> tuple[Tensor, Tensor]:
+        if context.ndim != 2 or context.shape[1] != self.in_features:
+            raise ValueError('Need padded number embeddings with token-validity masks')
+        slots = context.reshape(-1, MAX_NUMBER_TOKENS, self.width + 1)
+        return slots[..., :self.width], slots[..., self.width:]
+
+    def forward(self, context: Tensor) -> Tensor:
+        embeddings, mask = self.unpack(context.detach())
+        values = torch.nn.functional.linear(embeddings.double(), self.token_value.weight.double(),
+                                            self.token_value.bias.double())
+        state = values.new_zeros(len(values), 1)
+        for value, valid in zip(values.unbind(1), mask.unbind(1), strict=True):
+            state = torch.where(valid.bool(), self.gain.double() * state + value, state)
+        return state.to(self.token_value.weight.dtype)

scripts/recurrent_qwen/outputs.py

@@ -0,0 +1,56 @@
+"""Loop outputs and symbolic readout semantics."""
+
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor
+
+
+@dataclass
+class RecurrentOutput:
+    """Tuple index t-1 represents loop t; h_0 is not a decoded start symbol.
+
+    Logits are [batch, vocab] in answer mode or [batch, sequence, vocab]
+    in all-token mode. Optional hidden states are [batch, sequence, hidden].
+    Margins are [batch, loops]. Tensors retain autograd unless the caller
+    disables it; observations never feed back into the recurrent trajectory.
+    """
+
+    loop_logits: tuple[Tensor, ...]
+    hidden_states: tuple[Tensor, ...] | None = None
+    initial_hidden_state: Tensor | None = None
+    margins: Tensor | None = None
+    state_logits: Tensor | None = None  # Direct R readout [B,T,26], never fed back.
+    stop_logits: Tensor | None = None  # [batch, executed_loops], when completion is enabled.
+
+    @property
+    def logits(self) -> Tensor:
+        """The final loop's logits."""
+        return self.loop_logits[-1]
+
+
+def answer_margin(logits: Tensor, labels: Tensor, allowed_token_ids: Tensor) -> Tensor:
+    """Correct raw logit minus the best *other allowed* logit, per example.
+
+    Shapes: logits [batch, vocab], labels [batch] (token IDs), allowed IDs
+    [answers]. Labels must belong to a unique answer set of at least two IDs.
+    A tie has margin zero, so it is not margin-based correctness.
+    """
+    if logits.ndim != 2 or labels.shape != logits.shape[:1]:
+        raise ValueError("Expected logits [batch, vocab] and labels [batch]")
+    if labels.dtype != torch.long or allowed_token_ids.dtype != torch.long:
+        raise ValueError("Labels and allowed_token_ids must be torch.long token IDs")
+    if labels.device != logits.device or allowed_token_ids.device != logits.device:
+        raise ValueError("Logits, labels, and allowed_token_ids must share a device")
+    if allowed_token_ids.ndim != 1 or allowed_token_ids.numel() < 2:
+        raise ValueError("Provide at least two allowed token IDs in a 1D tensor")
+    if allowed_token_ids.unique().numel() != allowed_token_ids.numel():
+        raise ValueError("Allowed token IDs must be unique")
+    if (allowed_token_ids < 0).any() or (allowed_token_ids >= logits.shape[-1]).any():
+        raise ValueError("Allowed token IDs must be within the model vocabulary")
+    correct = labels[:, None] == allowed_token_ids[None, :]
+    if not correct.any(dim=-1).all():
+        raise ValueError("Every label must be an allowed token ID")
+    scores = logits.index_select(-1, allowed_token_ids)
+    wrong = scores.masked_fill(correct, -torch.inf).max(dim=-1).values
+    return logits.gather(-1, labels[:, None]).squeeze(-1) - wrong

scripts/recurrent_qwen/partitioned.py

@@ -0,0 +1,166 @@
+"""Inference-only causal factorization for count-dependent historical models.
+
+Identical prompt prefixes are computed once per graph, at *each* loop for
+full-sequence recurrence. Count-specific suffixes remain separate continuous
+states. Fixed-prompt recurrence retains first-pass layer K/V, replacing only
+the working position. No symbols, labels, depth values or stops enter R.
+"""
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+from transformers.models.qwen2.modeling_qwen2 import apply_rotary_pos_emb, eager_attention_forward
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+from .prefix import layer_compute
+
+
+def suffix_layer(layer: nn.Module, hidden: Tensor, embeddings: tuple[Tensor, Tensor],
+                 mask: Tensor, prefix: tuple[Tensor, Tensor],
+                 replace_positions: Tensor | None = None) -> tuple[Tensor, Tensor, Tensor]:
+    """Queries [B,Q,H], rotary K/V prefix [B,KV,P,D]; full additive query mask."""
+    attention = layer.self_attn
+    x = layer.input_layernorm(hidden)
+    shape = (*x.shape[:-1], -1, attention.head_dim)
+    q = attention.q_proj(x).view(shape).transpose(1, 2)
+    k = attention.k_proj(x).view(shape).transpose(1, 2)
+    v = attention.v_proj(x).view(shape).transpose(1, 2)
+    q, k = apply_rotary_pos_emb(q, k, *embeddings)
+    if replace_positions is None:
+        k, v = torch.cat((prefix[0], k), -2), torch.cat((prefix[1], v), -2)
+    else:
+        # Work on fresh tensors: original fixed memory is immutable across loops.
+        k_full, v_full = prefix[0].clone(), prefix[1].clone()
+        rows = torch.arange(len(hidden), device=hidden.device)
+        k_full[rows, :, replace_positions, :] = k[:, :, 0, :]
+        v_full[rows, :, replace_positions, :] = v[:, :, 0, :]
+        k, v = k_full, v_full
+    interface = ALL_ATTENTION_FUNCTIONS.get_interface(attention.config._attn_implementation, eager_attention_forward)
+    output, _ = interface(attention, q, k, v, mask,
+                         dropout=0.0, scaling=attention.scaling, sliding_window=None)
+    output = output.reshape(*x.shape[:-1], -1).contiguous()
+    result = hidden + attention.o_proj(output)
+    result = result + layer.mlp(layer.post_attention_layernorm(result))
+    return result, k, v
+
+
+@dataclass
+class PartitionedResult:
+    """Readouts [graph*count, loop, symbol], stop logits or None, final R working."""
+    scores: Tensor
+    stops: Tensor | None
+    working: Tensor
+
+
+@torch.inference_mode()
+def partitioned_forward(model: nn.Module, prefix_ids: Tensor, suffix_ids: Tensor,
+                        suffix_mask: Tensor, graph_indices: Tensor,
+                        token_ids: list[int], loops: int, horizons: Tensor | None = None) -> PartitionedResult:
+    """Exact causal factorization, modulo documented FP32 kernel roundoff.
+
+    Prefix [G,P] has no padding, suffix [B,S] is right padded. Each suffix maps
+    to one prefix via [B] graph_indices. Positions match the dense concatenation.
+    Supports historical full_sequence/fixed_prompt models, without routed inputs,
+    bridges, direct heads or private recurrent controllers. Never modifies model.
+    """
+    if (model.training or model.router is not None or model.bridge is not None or
+            model.state_head is not None or loops < 1):
+        raise ValueError('Partitioned path requires an unmodified historical eval model')
+    if prefix_ids.ndim != 2 or suffix_ids.shape != suffix_mask.shape or suffix_ids.ndim != 2:
+        raise ValueError('Expected prefix and aligned suffix token matrices')
+    if (not suffix_mask.any(-1).all() or not torch.equal(suffix_mask.long().cumprod(-1), suffix_mask.long())
+            or graph_indices.shape != (len(suffix_ids),) or
+            (graph_indices < 0).any() or (graph_indices >= len(prefix_ids)).any()):
+        raise ValueError('Need right-padded suffixes and valid graph indices')
+    prefix = model.embed_tokens(prefix_ids)
+    suffix = model.embed_tokens(suffix_ids)
+    device, dtype = prefix.device, prefix.dtype
+    p, s = prefix.shape[1], suffix.shape[1]
+    pp = torch.arange(p, device=device)[None]
+    sp = torch.arange(p, p + s, device=device)[None]
+    pe, se = model.rotary_emb(prefix, pp), model.rotary_emb(suffix, sp)
+    negative = torch.finfo(dtype).min
+    pm = torch.zeros(p, p, device=device, dtype=dtype).masked_fill(
+        torch.arange(p, device=device)[None] > torch.arange(p, device=device)[:, None], negative)[None, None]
+    keep = torch.cat((torch.ones(len(suffix), p, device=device, dtype=torch.bool), suffix_mask.bool()), -1)
+    allowed = (torch.arange(p + s, device=device)[None] <= sp[0, :, None])[None] & keep[:, None, :]
+    sm = torch.zeros_like(allowed, dtype=dtype).masked_fill(~allowed, negative)[:, None]
+    ends = suffix_mask.sum(-1).long() - 1
+    rows = torch.arange(len(suffix), device=device)
+    am = sm[rows, :, ends, :].unsqueeze(2)
+    ae = tuple(e.expand(len(suffix), -1, -1)[rows, ends, :].unsqueeze(1) for e in se)
+
+    def block(layers, ph, sh, record=False):
+        memory = []
+        for layer in layers:
+            ph, k, v = layer_compute(layer, ph, pe, pm)
+            sh, keys, values = suffix_layer(layer, sh, se, sm,
+                                           (k.index_select(0, graph_indices), v.index_select(0, graph_indices)))
+            if record:
+                memory.append((keys, values))
+        return ph, sh, memory
+
+    def working_block(layers, state, memory):
+        for layer, kv in zip(layers, memory, strict=True):
+            state, _, _ = suffix_layer(layer, state, ae, am, kv, ends + p)
+        return state
+
+    prefix, suffix, _ = block(model.prelude.layers, prefix, suffix)
+    fixed = model.recurrence_mode == 'fixed_prompt'
+    scores, stops = [], []
+    original_batch = len(suffix_ids)
+    active_rows = torch.arange(original_batch, device=device)
+    crossed = torch.zeros(original_batch, device=device, dtype=torch.bool)
+    if horizons is not None and (horizons.shape != (original_batch,) or (horizons < 1).any() or (horizons > loops).any()):
+        raise ValueError('Requested readout horizons must fit the loop cap')
+    r_memory = c_memory = None
+    working = None
+    selected = model.lm_head.weight[token_ids]
+    for loop in range(loops):
+        if fixed and loop:
+            working = working_block(model.recurrent.layers, working, r_memory)
+            decoded = working_block(model.coda.layers, working, c_memory)
+        else:
+            prefix, suffix, r_memory = block(model.recurrent.layers, prefix, suffix, fixed)
+            working = suffix[rows, ends, :].unsqueeze(1)
+            _, decoded_suffix, c_memory = block(model.coda.layers, prefix, suffix, fixed)
+            decoded = decoded_suffix[rows, ends, :].unsqueeze(1)
+        score = nn.functional.linear(model.norm(decoded[:, 0]), selected).float()
+        if not torch.isfinite(score).all():
+            raise FloatingPointError('Nonfinite partitioned readout')
+        full_score = score.new_zeros(original_batch, score.shape[-1])
+        full_score[active_rows] = score
+        scores.append(full_score)
+        if model.completion_head is not None:
+            stop = model.completion_head(working[:, 0]).float()
+            if not torch.isfinite(stop).all():
+                raise FloatingPointError('Nonfinite partitioned completion head')
+            full_stop = stop.new_full((original_batch,), float('-inf'))
+            full_stop[active_rows] = stop
+            stops.append(full_stop)
+            crossed[active_rows] |= stop >= 0
+        if horizons is not None and loop + 1 < loops:
+            need = horizons[active_rows] > loop + 1
+            if model.completion_head is not None:
+                need |= ~crossed[active_rows]
+            # A CUDA boolean index discovers its output shape on the host.
+            # Discover the same ordered rows once, then reuse them for every
+            # state/cache tensor instead of repeating that synchronization.
+            kept = need.nonzero().flatten()
+            if not kept.numel():
+                for _ in range(loop + 1, loops):
+                    scores.append(full_score.new_zeros(full_score.shape))
+                    if stops: stops.append(full_stop.new_full(full_stop.shape, float('-inf')))
+                break
+            active_rows = active_rows.index_select(0, kept)
+            graph_indices = graph_indices.index_select(0, kept)
+            suffix = suffix.index_select(0, kept)
+            working = working.index_select(0, kept)
+            ends = ends.index_select(0, kept)
+            rows = torch.arange(len(active_rows), device=device)
+            sm = sm.index_select(0, kept); am = am.index_select(0, kept)
+            ae = tuple(e.index_select(0, kept) for e in ae)
+            if fixed:
+                r_memory = [(key.index_select(0, kept), value.index_select(0, kept)) for key, value in r_memory]
+                c_memory = [(key.index_select(0, kept), value.index_select(0, kept)) for key, value in c_memory]
+    return PartitionedResult(torch.stack(scores, 1), torch.stack(stops, 1) if stops else None, working[:, 0])

scripts/recurrent_qwen/prefix.py

@@ -0,0 +1,67 @@
+"""Differentiable fixed-prefix execution; one working query after the first pass."""
+import torch
+from torch import Tensor, nn
+from .memory import PromptMemory
+from torch.utils.checkpoint import checkpoint
+from transformers.models.qwen2.modeling_qwen2 import apply_rotary_pos_emb, eager_attention_forward
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+
+def pick(x: Tensor, positions: list[int]) -> Tensor:
+    return torch.stack([x[row, position:position + 1] for row, position in enumerate(positions)])
+
+
+def layer_compute(layer: nn.Module, hidden: Tensor, embeddings: tuple[Tensor, Tensor],
+                  mask: Tensor | None, positions: list[int] | None = None,
+                  cached: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, Tensor, Tensor]:
+    """Pure operation so non-reentrant checkpointing preserves prefix gradients."""
+    attention = layer.self_attn
+    residual = hidden
+    x = layer.input_layernorm(hidden)
+    shape = (*x.shape[:-1], -1, attention.head_dim)
+    q = attention.q_proj(x).view(shape).transpose(1, 2)
+    k = attention.k_proj(x).view(shape).transpose(1, 2)
+    v = attention.v_proj(x).view(shape).transpose(1, 2)
+    q, k = apply_rotary_pos_emb(q, k, *embeddings)
+    if cached is not None:
+        keys, values = cached
+        active = (torch.arange(keys.shape[-2], device=x.device)[None, :] ==
+                  torch.tensor(positions, device=x.device)[:, None])[:, None, :, None]
+        k = torch.where(active, k, keys)
+        v = torch.where(active, v, values)
+    interface = ALL_ATTENTION_FUNCTIONS.get_interface(attention.config._attn_implementation, eager_attention_forward)
+    output, _ = interface(attention, q, k, v, mask,
+        dropout=attention.attention_dropout if layer.training else 0.0,
+        scaling=attention.scaling, sliding_window=None)
+    output = output.reshape(*x.shape[:-1], -1).contiguous()
+    hidden = residual + attention.o_proj(output)
+    hidden = hidden + layer.mlp(layer.post_attention_layernorm(hidden))
+    return hidden, k, v
+
+
+def fixed_prefix_layer(layer: nn.Module, hidden: Tensor, embeddings: tuple[Tensor, Tensor],
+                       mask: Tensor | None, memory: PromptMemory, index: int,
+                       use_checkpoint: bool) -> Tensor:
+    """Reuse prefix K/V without detaching it or caching across model forwards."""
+    positions = memory.write_mask[..., 0].long().argmax(-1).tolist()
+    cached = None
+    first = memory.output is None
+    if not first:
+        hidden = pick(hidden, positions)
+        embeddings = tuple(pick(e.expand(memory.write_mask.shape[0], -1, -1), positions) for e in embeddings)
+        cached = memory.keys_values[index]
+        if mask is None:
+            keep = torch.arange(cached[0].shape[-2], device=hidden.device)[None, :] <= torch.tensor(positions, device=hidden.device)[:, None]
+            mask = torch.zeros_like(keep, dtype=hidden.dtype).masked_fill(~keep, torch.finfo(hidden.dtype).min)[:, None, None, :]
+        else:
+            mask = torch.stack([mask[row if mask.shape[0] > 1 else 0, :, pos:pos + 1, :] for row, pos in enumerate(positions)])
+    def compute(x, cos, sin, *kv):
+        return layer_compute(layer, x, (cos, sin), mask, positions, tuple(kv) if kv else None)
+    args = (hidden, *embeddings, *(cached or ()))
+    result, keys, values = (checkpoint(compute, *args, use_reentrant=False)
+                            if use_checkpoint and torch.is_grad_enabled() else compute(*args))
+    if first:
+        memory.keys_values.append((keys, values))
+        memory.layer_outputs.append(result)
+        return result
+    return torch.where(memory.write_mask, result, memory.layer_outputs[index])

scripts/recurrent_qwen/validation.py

@@ -0,0 +1,185 @@
+"""Stage 0 checks shared by the pretrained CLI and a small integration test."""
+
+from collections.abc import Callable
+from time import perf_counter
+from typing import Any
+
+import torch
+from torch import Tensor
+from torch.nn import functional as F
+from transformers import Qwen2ForCausalLM
+
+from .lora_utils import attach_recurrent_lora
+from .model import RecurrentQwen
+
+
+def validate_stage0(
+    base: Qwen2ForCausalLM,
+    cases: dict[str, dict[str, Tensor]],
+    *,
+    recurrent_start: int = 6,
+    recurrent_end: int = 18,
+    rank: int = 8,
+    alpha: int = 16,
+    loop_counts: tuple[int, ...] = (1, 2, 4),
+    atol: float = 1e-5,
+    rtol: float = 1e-5,
+    progress: Callable[[str], None] = print,
+) -> dict[str, Any]:
+    """Assert architecture contracts; return serializable measured evidence.
+
+    Mutates/freezes base and attaches fresh LoRA. No optimizer step is taken.
+    The supplied cases must be short, nonempty input_ids/attention_mask pairs;
+    equivalence compares every vocabulary logit at every sequence position.
+    """
+    if not cases or not loop_counts or min(loop_counts) < 1 or max(loop_counts) < 2:
+        raise ValueError("Provide input cases and positive loop counts including a depth > 1")
+    base.eval()
+    base.config.use_cache = False
+    originals = tuple(base.parameters())
+    original_ids = {id(p) for p in originals}
+    original_layers = tuple(base.model.layers)
+    progress("Computing ordinary-Qwen reference logits")
+    with torch.no_grad():
+        references = {name: base(**inputs, use_cache=False).logits for name, inputs in cases.items()}
+    model = RecurrentQwen(base, recurrent_start, recurrent_end).eval()
+    assert tuple(model.recurrent.layers) == original_layers[recurrent_start:recurrent_end]
+    errors = []
+    for adapted in (False, True):
+        if adapted:
+            attach_recurrent_lora(model, rank=rank, alpha=alpha)
+        progress(f"Checking T=1 equivalence ({'with' if adapted else 'without'} LoRA)")
+        with torch.no_grad():
+            for name, inputs in cases.items():
+                actual = model(**inputs, logits_mode="all").logits
+                expected = references[name]
+                torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+                difference = (actual - expected).abs()
+                errors.append({
+                    "case": name,
+                    "lora": adapted,
+                    "max_absolute_error": difference.max().item(),
+                    "mean_absolute_error": difference.mean().item(),
+                    "logits_shape": list(actual.shape),
+                })
+    del references
+
+    progress("Checking shared weights, loop observability, and configurable depth")
+    parameters_before = tuple(id(p) for p in model.parameters())
+    calls: list[tuple[int, tuple[int, ...], Tensor, Tensor]] = []
+    block_calls: list[str] = []
+    layer_calls: list[int] = []
+    handles = [model.recurrent.register_forward_hook(
+        lambda module, args, output: calls.append(
+            (id(module), tuple(id(p) for p in module.parameters()), args[0], output)
+        )
+    )]
+    for name in ("prelude", "coda"):
+        handles.append(getattr(model, name).register_forward_hook(
+            lambda module, args, output, name=name: block_calls.append(name)
+        ))
+    for layer in model.recurrent.layers:
+        handles.append(layer.register_forward_hook(
+            lambda module, args, output: layer_calls.append(id(module))
+        ))
+    inputs = next(iter(cases.values()))
+    batch, sequence = inputs["input_ids"].shape
+    device = inputs["input_ids"].device
+    allowed = inputs["input_ids"].unique()
+    if allowed.numel() < 2:
+        raise ValueError("Validation inputs must contain at least two distinct token IDs")
+    recurrent_ids = tuple(id(p) for p in model.recurrent.parameters())
+    trajectories = []
+    try:
+        with torch.no_grad():
+            for depth in loop_counts:
+                calls.clear()
+                block_calls.clear()
+                layer_calls.clear()
+                labels = allowed[torch.arange(depth, device=device) % allowed.numel()].expand(batch, -1)
+                start = perf_counter()
+                output = model(**inputs, num_loops=depth, return_hidden_states=True,
+                               labels=labels, allowed_token_ids=allowed)
+                if device.type == "mps":
+                    torch.mps.synchronize()
+                seconds = perf_counter() - start
+                assert len(calls) == depth
+                assert block_calls.count("prelude") == 1 and block_calls.count("coda") == depth
+                assert layer_calls == [id(layer) for layer in model.recurrent.layers] * depth
+                assert all(call[0] == id(model.recurrent) and call[1] == recurrent_ids for call in calls)
+                assert tuple(id(p) for p in model.parameters()) == parameters_before
+                assert calls[0][2] is output.initial_hidden_state
+                assert all(calls[t][2] is calls[t - 1][3] for t in range(1, depth))
+                assert len(output.hidden_states) == len(output.loop_logits) == depth
+                assert output.margins.shape == (batch, depth)
+                for hidden, logits in zip(output.hidden_states, output.loop_logits):
+                    assert hidden.shape == (batch, sequence, model.config.hidden_size)
+                    assert logits.shape == (batch, model.config.vocab_size)
+                    assert torch.isfinite(hidden).all() and torch.isfinite(logits).all()
+                assert torch.isfinite(output.margins).all()
+                trajectories.append({
+                    "num_loops": depth,
+                    "seconds": seconds,
+                    "recurrent_calls": len(calls),
+                    "recurrent_layer_calls": len(layer_calls),
+                    "prelude_calls": block_calls.count("prelude"),
+                    "coda_calls": block_calls.count("coda"),
+                    "predicted_token_ids": [scores.argmax(-1).tolist() for scores in output.loop_logits],
+                    "diagnostic_labels": labels.tolist(),
+                    "allowed_token_ids": allowed.tolist(),
+                    "margins": output.margins.tolist(),
+                })
+    finally:
+        for handle in handles:
+            handle.remove()
+    calls.clear()
+    del output
+
+    progress("Checking two-loop backward through the frozen coda (no optimizer step)")
+    model.train()
+    model.zero_grad(set_to_none=True)
+    start = perf_counter()
+    output = model(**inputs, num_loops=2, return_hidden_states=True)
+    for hidden in output.hidden_states:
+        hidden.retain_grad()
+    # Arbitrary token labels test graph connectivity, not task competence.
+    loss = F.cross_entropy(output.logits.float(), inputs["input_ids"][:, 0])
+    loss.backward()
+    if device.type == "mps":
+        torch.mps.synchronize()
+    backward_seconds = perf_counter() - start
+    assert all(not p.requires_grad and p.grad is None for p in originals)
+    trainable = [(name, p) for name, p in model.named_parameters() if p.requires_grad]
+    assert len(trainable) == (recurrent_end - recurrent_start) * 2 * 2
+    assert {id(p) for p in model.parameters()} - original_ids == {id(p) for _, p in trainable}
+    gradients = {}
+    for name, parameter in trainable:
+        assert name.startswith("recurrent.") and (".lora_A." in name or ".lora_B." in name)
+        assert parameter.grad is not None and torch.isfinite(parameter.grad).all(), name
+        norm = parameter.grad.float().norm().item()
+        if ".lora_B." in name:
+            assert norm > 0, name
+        gradients[name] = {"elements": parameter.numel(), "gradient_l2": norm}
+    hidden_gradients = [hidden.grad.float().norm().item() for hidden in output.hidden_states]
+    assert all(value > 0 and torch.isfinite(torch.tensor(value)) for value in hidden_gradients)
+    result = {
+        "gate_passed": True,
+        "checks": {name: "PASS" for name in (
+            "t1_equivalence", "shared_weights", "gradient_scope", "observability", "configurable_depth"
+        )},
+        "split": {"prelude": [0, recurrent_start], "recurrent": [recurrent_start, recurrent_end],
+                  "coda": [recurrent_end, len(original_layers)], "end_indices_exclusive": True},
+        "lora": {"rank": rank, "alpha": alpha, "targets": ["q_proj", "v_proj"], "dropout": 0.0},
+        "bridge": None,
+        "tolerance": {"atol": atol, "rtol": rtol},
+        "equivalence": errors,
+        "trajectories": trajectories,
+        "parameters": {"original": sum(p.numel() for p in originals),
+                       "trainable": sum(p.numel() for _, p in trainable),
+                       "trainable_tensors": len(trainable)},
+        "gradient_check": {"loss": loss.item(), "forward_backward_seconds": backward_seconds,
+                           "hidden_gradient_l2": hidden_gradients, "parameters": gradients,
+                           "frozen_parameters_with_grad": 0, "optimizer_steps": 0},
+    }
+    model.zero_grad(set_to_none=True)
+    return result

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3fd169731d2cbde95e10bf356d66d5997fd885dd8dbb6fb4684da3f23b2585d8
+size 11421892

tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "is_local": false,
+  "local_files_only": true,
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}