Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,12 @@ 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
+assets/architecture.png filter=lfs diff=lfs merge=lfs -text
+assets/reasoning_loop.png filter=lfs diff=lfs merge=lfs -text
+assets/vram_comparison.png filter=lfs diff=lfs merge=lfs -text
+default/d/40b6413d8d045553235110cf8a3113dc filter=lfs diff=lfs merge=lfs -text
+default/ocdbt.process_0/d/2655f27744aa28bc57a54732ca8aa17f filter=lfs diff=lfs merge=lfs -text
+default/ocdbt.process_0/d/309bfd1f96632d6760dd55dea979babf filter=lfs diff=lfs merge=lfs -text
+default/ocdbt.process_0/d/37b57931b1bb0df657d81dd245946279 filter=lfs diff=lfs merge=lfs -text
+default/ocdbt.process_0/d/7d3a8dd28172f4fc4fe186eaa73f2843 filter=lfs diff=lfs merge=lfs -text
+default/ocdbt.process_0/d/cca1e2cb6509e2bae33156603f3ff2de filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,306 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+language:
+  - en
+tags:
+  - jax
+  - flax
+  - language-model
+  - text-generation
+  - retrieval-augmented
+  - custom-architecture
+  - research
+library_name: flax
+pipeline_tag: text-generation
+model_type: dpsnr
+datasets:
+  - fineweb
+metrics:
+  - perplexity
+inference: false
+widget:
+  - text: "The future of artificial intelligence"
+    example_title: "AI Future"
+  - text: "Once upon a time in a land"
+    example_title: "Story"
+  - text: "The key insight of this paper is"
+    example_title: "Research"
+model-index:
+  - name: DPSNR-Large
+    results: []
+---
+
+# DPSNR — Dynamic Parameter Selection Network with Reasoning
+
+> **A JAX/Flax language model that separates *what it knows* from *how it thinks* — so the knowledge can grow to 100B+ vectors while inference stays fast and cheap.**
+
+---
+
+## What Is DPSNR?
+
+Normal large language models (GPT, Llama, etc.) mix logic and facts together inside the same transformer weights. When you want more knowledge, you need more parameters, which means more GPU VRAM, more compute, more cost — the **VRAM Wall**.
+
+DPSNR breaks that wall. It splits the model into two parts:
+
+| Part | Role | Size |
+|------|------|------|
+| **TinyController** | Does the thinking / reasoning | ~350M params on GPU |
+| **CoordinateMassivePool** | Stores world knowledge as vectors | 262K–1T+ vectors, can live on disk |
+
+The controller *queries* the pool each reasoning step instead of storing facts in its weights. Pool size can grow arbitrarily; inference cost stays **O(1)**.
+
+---
+
+## Architecture Overview
+
+![DPSNR Architecture](assets/architecture.png)
+
+The model has **4 components** that work together:
+
+```mermaid
+flowchart LR
+    Input["🗒️ Input Tokens"] --> TC
+
+    subgraph TC["① TinyController"]
+        direction TB
+        E["Token + Position\nEmbedding"] --> TL["12× Transformer\nLayers (768-dim)"] --> H["Hidden States\n(B, T, 768)"]
+    end
+
+    TC --> LI
+
+    subgraph LI["② LearnedIndexer"]
+        direction TB
+        AP["Attention Pooling\n(learn which token to query from)"] --> MH["Multi-Head Dense\n→ μ coordinate\n→ σ bandwidth"]
+    end
+
+    LI -->|"μ, σ"| Pool
+
+    subgraph Pool["③ CoordinateMassivePool"]
+        direction TB
+        PS["262,144 × 768\nlearned vectors"] --> GW["Gaussian window\naround μ ± K vectors\nweighted by σ"] --> AV["Aggregated\nKnowledge Vector\n(B, 768)"]
+    end
+
+    Pool --> ACC
+
+    subgraph ACC["④ Adaptive Compute Controller"]
+        direction TB
+        RI["Integrate knowledge\ninto hidden state"] --> HN["Halt Network\n(should we stop?)"]
+        HN -->|"halt < 0.99"| RI
+    end
+
+    ACC -->|"Final hidden state"| Out["📝 Output Logits\n(B, T, vocab)"]
+    ACC -->|"loop back\n(up to 6 times)"| TC
+```
+
+---
+
+## How the Reasoning Loop Works
+
+Instead of doing one pass like most LLMs, DPSNR thinks iteratively — like a human reading and re-reading a hard problem.
+
+![Reasoning Loop](assets/reasoning_loop.png)
+
+Each loop:
+1. **TinyController** encodes the input → produces a hidden state
+2. **LearnedIndexer** converts the hidden state into a *coordinate* (μ) and *uncertainty* (σ)
+3. **CoordinateMassivePool** retrieves K=32 knowledge vectors near μ, weighted by a Gaussian of width σ
+4. Retrieved knowledge is fused into the hidden state
+5. **ACC** decides: confident enough? → output. Unsure? → loop again
+
+Simple questions finish in 1–2 loops. Hard questions use all 6. Compute is spent where it's needed.
+
+---
+
+## Breaking the VRAM Wall
+
+![VRAM Comparison](assets/vram_comparison.png)
+
+A 70B dense model requires 80GB+ of expensive HBM VRAM just to load. Because DPSNR stores knowledge as a flat array of vectors (not entangled with transformer weights), the pool can live in:
+
+- **System RAM** — 64GB RAM holds ~130M vectors × 768-dim at float32
+- **NVMe SSD** — mmap'd; only the retrieved window is paged in
+- **GPU VRAM** — only the TinyController (~1.3GB at bf16) needs the GPU
+
+```
+Dense 70B:  [GPU|▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓ 80GB VRAM ▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓]
+DPSNR:      [GPU|▓ 4GB]  +  [RAM|▓▓▓▓▓▓ Pool ▓▓▓▓▓▓]  ← no problem
+```
+
+---
+
+## Quick Start (Inference)
+
+### 1. Activate the virtualenv
+
+```bash
+cd /path/to/dpsn
+source .venv/bin/activate
+```
+
+### 2. Verify GPU is available
+
+```bash
+python -c "import jax; print(jax.devices())"
+# → [CudaDevice(id=0)]
+```
+
+### 3. Run inference
+
+```bash
+# Single prompt
+python infer.py --prompt "The future of artificial intelligence"
+
+# Interactive chat mode
+python infer.py
+
+# All options
+python infer.py \
+    --prompt    "Once upon a time" \
+    --max_tokens 200 \
+    --temp       0.8 \
+    --top_k      50 \
+    --penalty    1.3
+```
+
+The first run takes ~20–30s to JIT-compile the forward pass. Subsequent prompts in the same session are fast.
+
+---
+
+## Inference Script — `infer.py`
+
+The file `infer.py` is **fully self-contained** — it has the entire model architecture, checkpoint loading, and generation logic in one file. No dependency on the `dpsn_r_jax` package.
+
+```
+infer.py
+├── DPSNRConfig          ← Large model config, hardcoded
+├── FlashCausalSelfAttention
+├── TinyFFN / TinyTransformerLayer
+├── TinyController       ← 12-layer transformer encoder + LM head
+├── LearnedIndexer       ← μ, σ coordinate predictor
+├── CoordinateMassivePool  ← 1D flat pool (used by large config)
+├── CoordinateMassivePool2D ← 2D grid pool (use_2d_pool=True)
+├── AdaptiveComputeController ← halt/loop decision
+├── DPSNR                ← full forward pass, reasoning scan
+├── TrainState           ← pytree-compatible state for orbax restore
+├── load_checkpoint()    ← restores params only (no optimizer bloat)
+├── _forward()           ← @jax.jit compiled forward pass
+└── generate()           ← autoregressive sampling, fixed-size buffers
+```
+
+### CLI arguments
+
+| Argument | Default | Description |
+|---|---|---|
+| `--prompt` | None | Text prompt. Omit to enter interactive mode |
+| `--max_tokens` | 100 | Maximum new tokens to generate |
+| `--temp` | 0.7 | Sampling temperature. Lower = more focused |
+| `--top_k` | 40 | Only sample from top-K most likely tokens |
+| `--penalty` | 1.2 | Repetition penalty. >1 discourages repeats |
+| `--checkpoint_dir` | `./checkpoints_dir` | Override checkpoint path |
+
+---
+
+## Model Configuration (Large)
+
+The `large` config is hardcoded in `infer.py`:
+
+```python
+DPSNRConfig(
+    vocab_size            = 50257,   # GPT-Neo tokenizer vocab
+    controller_hidden_dim = 768,     # transformer width
+    controller_num_layers = 12,      # transformer depth
+    controller_num_heads  = 12,      # attention heads
+    max_seq_len           = 1024,    # max context window
+    pool_total_vectors    = 262144,  # 2^18 knowledge vectors
+    pool_hidden_dim       = 768,     # vector dimension
+    max_reasoning_loops   = 6,       # max iterations of the loop
+)
+```
+
+### Model size breakdown
+
+```mermaid
+pie title DPSNR Large — Parameter Distribution (~350M total)
+    "CoordinateMassivePool (262K × 768)" : 201
+    "TinyController (12L × 768d)" : 85
+    "LearnedIndexer" : 3
+    "AdaptiveComputeController" : 2
+    "Retrieval Integrator" : 9
+```
+
+---
+
+## Tokenizer
+
+Uses **`EleutherAI/gpt-neo-125M`** tokenizer — GPT-2 compatible BPE with 50,257 tokens. Downloaded automatically via HuggingFace on first use.
+
+---
+
+
+## Key Ideas Explained Simply
+
+### Why the pool doesn't slow things down
+
+Every retrieval fetches exactly `K=32` vectors regardless of pool size. Going from 10K to 100B pool vectors doesn't add a single FLOP — only the storage grows.
+
+```mermaid
+xychart-beta
+    title "Inference Latency vs Pool Size"
+    x-axis ["10K vectors", "100K", "262K", "1M", "1B", "100B"]
+    y-axis "Relative Latency" 0 --> 2
+    line [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
+```
+
+### Why Gaussian retrieval is better than nearest-neighbour
+
+Nearest-neighbour lookup (like a typical vector database) must search the entire pool. DPSNR uses a **coordinate** approach: the pool is arranged in a continuous 1D (or 2D grid) space. The indexer predicts a *position* μ and *width* σ, and we simply slice a window. No search required — it's a direct lookup with `jax.lax.dynamic_slice`.
+
+### Why σ matters
+
+- **Small σ** → sharp, precise retrieval (good for exact facts, code syntax)
+- **Large σ** → broad, averaged retrieval (good for general context)
+
+σ is learned per token, per reasoning step — the model naturally figures out how precise to be.
+
+---
+
+## Performance
+
+| Metric | Value | Notes |
+|---|---|---|
+| Training platform | TPU v5e-8 | 8-chip pod slice |
+| Throughput | **240–250K tokens/sec** | HBM bandwidth bound |
+| Sustained compute | **260–270 TFLOPS** | Below 393 TFLOPS peak |
+| Bottleneck | Memory bandwidth | Pool gather ops, not MXU |
+| Optimizer speedup vs dense | **590×** | Sparse Adam on retrieved indices only |
+| Checkpoint step | 31,000 | |
+| GPU VRAM (inference) | ~1.3GB (params only, bf16) | Pool can live off-device |
+| Inference tested on | NVIDIA RTX 2050 (4GB) | Consumer GPU confirmed |
+
+---
+
+## Dependencies
+
+```
+jax + jaxlib   ← Core ML framework (GPU/TPU backend)
+flax           ← Neural network layers and module API
+optax          ← Optimizers (used for checkpoint structure only)
+orbax          ← Checkpoint save/restore
+transformers   ← Tokenizer (HuggingFace)
+```
+
+Install:
+```bash
+pip install jax jaxlib flax optax orbax-checkpoint transformers
+```
+
+---
+
+## Citation / Reference
+
+```
+DPSNR: Disaggregated Parameter Selection Network with Reasoning
+Architecture: TinyController + CoordinateMassivePool + LearnedIndexer + ACC
+Implementation: JAX/Flax
+Checkpoint: step 31,000
+```

_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": {"default": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1772564676365471524, "commit_timestamp_nsecs": 1772564878381366150, "custom_metadata": {}}

default/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "step", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.acc.halt_net.layers_0.bias", "write_shape": [24], "chunk_shape": [24], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.acc.halt_net.layers_0.kernel", "write_shape": [96, 192], "chunk_shape": [96, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.acc.halt_net.layers_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.acc.halt_net.layers_2.kernel", "write_shape": [24, 1], "chunk_shape": [24, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.acc.loop_embed.embedding", "write_shape": [4, 768], "chunk_shape": [4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.acc.state_gate.layers_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.acc.state_gate.layers_0.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.acc.state_norm.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.acc.state_norm.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.acc.state_transform.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.acc.state_transform.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.embedding.embedding", "write_shape": [50257, 96], "chunk_shape": [50257, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.final_norm.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.final_norm.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_0.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_0.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_0.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_0.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_0.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_0.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_0.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_0.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_0.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_0.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_1.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_1.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_1.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_1.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_1.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_1.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_1.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_1.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_1.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_1.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_10.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_10.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_10.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_10.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_10.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_10.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_10.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_10.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_10.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_10.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_11.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_11.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_11.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_11.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_11.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_11.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_11.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_11.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_11.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_11.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_2.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_2.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_2.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_2.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_2.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_2.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_2.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_2.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_2.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_2.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_3.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_3.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_3.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_3.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_3.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_3.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_3.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_3.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_3.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_3.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_4.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_4.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_4.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_4.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_4.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_4.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_4.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_4.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_4.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_4.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_5.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_5.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_5.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_5.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_5.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_5.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_5.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_5.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_5.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_5.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_6.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_6.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_6.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_6.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_6.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_6.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_6.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_6.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_6.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_6.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_7.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_7.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_7.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_7.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_7.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_7.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_7.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_7.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_7.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_7.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_8.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_8.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_8.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_8.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_8.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_8.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_8.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_8.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_8.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_8.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_9.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_9.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_9.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_9.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_9.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_9.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_9.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_9.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_9.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.layers_9.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.lm_head.kernel", "write_shape": [96, 50257], "chunk_shape": [96, 50257], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.controller.pos_encoding.embedding", "write_shape": [128, 768], "chunk_shape": [128, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.indexer.Dense_0.kernel", "write_shape": [96, 1], "chunk_shape": [96, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.indexer.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.indexer.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.indexer.Dense_2.bias", "write_shape": [48], "chunk_shape": [48], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.indexer.Dense_2.kernel", "write_shape": [96, 384], "chunk_shape": [96, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.indexer.Dense_3.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.indexer.Dense_3.kernel", "write_shape": [48, 1], "chunk_shape": [48, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.indexer.Dense_4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.indexer.Dense_4.kernel", "write_shape": [48, 1], "chunk_shape": [48, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.pool.params_storage", "write_shape": [32768, 768], "chunk_shape": [32768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.retrieval_integrator.layers_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.retrieval_integrator.layers_0.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.retrieval_integrator.layers_2.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.retrieval_integrator.layers_2.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.retrieval_integrator.layers_3.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.retrieval_integrator.layers_3.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.acc.halt_net.layers_0.bias", "write_shape": [24], "chunk_shape": [24], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.acc.halt_net.layers_0.kernel", "write_shape": [96, 192], "chunk_shape": [96, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.acc.halt_net.layers_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.acc.halt_net.layers_2.kernel", "write_shape": [24, 1], "chunk_shape": [24, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.acc.loop_embed.embedding", "write_shape": [4, 768], "chunk_shape": [4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.acc.state_gate.layers_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.acc.state_gate.layers_0.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.acc.state_norm.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.acc.state_norm.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.acc.state_transform.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.acc.state_transform.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.embedding.embedding", "write_shape": [50257, 96], "chunk_shape": [50257, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.final_norm.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.final_norm.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_0.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_0.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_0.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_0.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_0.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_0.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_0.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_0.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_0.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_0.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_1.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_1.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_1.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_1.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_1.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_1.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_1.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_1.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_1.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_1.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_10.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_10.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_10.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_10.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_10.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_10.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_10.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_10.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_10.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_10.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_11.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_11.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_11.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_11.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_11.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_11.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_11.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_11.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_11.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_11.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_2.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_2.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_2.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_2.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_2.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_2.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_2.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_2.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_2.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_2.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_3.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_3.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_3.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_3.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_3.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_3.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_3.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_3.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_3.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_3.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_4.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_4.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_4.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_4.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_4.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_4.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_4.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_4.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_4.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_4.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_5.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_5.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_5.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_5.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_5.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_5.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_5.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_5.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_5.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_5.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_6.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_6.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_6.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_6.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_6.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_6.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_6.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_6.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_6.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_6.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_7.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_7.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_7.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_7.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_7.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_7.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_7.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_7.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_7.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_7.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_8.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_8.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_8.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_8.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_8.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_8.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_8.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_8.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_8.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_8.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_9.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_9.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_9.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_9.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_9.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_9.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_9.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_9.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_9.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.layers_9.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.lm_head.kernel", "write_shape": [96, 50257], "chunk_shape": [96, 50257], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.controller.pos_encoding.embedding", "write_shape": [128, 768], "chunk_shape": [128, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.indexer.Dense_0.kernel", "write_shape": [96, 1], "chunk_shape": [96, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.indexer.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.indexer.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.indexer.Dense_2.bias", "write_shape": [48], "chunk_shape": [48], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.indexer.Dense_2.kernel", "write_shape": [96, 384], "chunk_shape": [96, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.indexer.Dense_3.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.indexer.Dense_3.kernel", "write_shape": [48, 1], "chunk_shape": [48, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.indexer.Dense_4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.indexer.Dense_4.kernel", "write_shape": [48, 1], "chunk_shape": [48, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.retrieval_integrator.layers_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.retrieval_integrator.layers_0.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.retrieval_integrator.layers_2.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.retrieval_integrator.layers_2.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.retrieval_integrator.layers_3.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.retrieval_integrator.layers_3.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.acc.halt_net.layers_0.bias", "write_shape": [24], "chunk_shape": [24], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.acc.halt_net.layers_0.kernel", "write_shape": [96, 192], "chunk_shape": [96, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.acc.halt_net.layers_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.acc.halt_net.layers_2.kernel", "write_shape": [24, 1], "chunk_shape": [24, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.acc.loop_embed.embedding", "write_shape": [4, 768], "chunk_shape": [4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.acc.state_gate.layers_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.acc.state_gate.layers_0.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.acc.state_norm.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.acc.state_norm.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.acc.state_transform.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.acc.state_transform.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.embedding.embedding", "write_shape": [50257, 96], "chunk_shape": [50257, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.final_norm.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.final_norm.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_0.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_0.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_0.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_0.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_0.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_0.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_0.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_0.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_0.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_0.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_1.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_1.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_1.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_1.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_1.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_1.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_1.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_1.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_1.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_1.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_10.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_10.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_10.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_10.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_10.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_10.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_10.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_10.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_10.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_10.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_11.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_11.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_11.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_11.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_11.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_11.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_11.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_11.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_11.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_11.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_2.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_2.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_2.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_2.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_2.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_2.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_2.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_2.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_2.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_2.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_3.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_3.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_3.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_3.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_3.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_3.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_3.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_3.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_3.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_3.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_4.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_4.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_4.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_4.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_4.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_4.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_4.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_4.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_4.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_4.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_5.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_5.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_5.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_5.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_5.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_5.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_5.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_5.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_5.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_5.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_6.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_6.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_6.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_6.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_6.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_6.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_6.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_6.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_6.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_6.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_7.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_7.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_7.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_7.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_7.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_7.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_7.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_7.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_7.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_7.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_8.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_8.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_8.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_8.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_8.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_8.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_8.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_8.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_8.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_8.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_9.FlashCausalSelfAttention_0.Dense_0.kernel", "write_shape": [96, 2304], "chunk_shape": [96, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_9.FlashCausalSelfAttention_0.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_9.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_9.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_9.LayerNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_9.LayerNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_9.TinyFFN_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_9.TinyFFN_0.Dense_0.kernel", "write_shape": [96, 1536], "chunk_shape": [96, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_9.TinyFFN_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.layers_9.TinyFFN_0.Dense_1.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.lm_head.kernel", "write_shape": [96, 50257], "chunk_shape": [96, 50257], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.controller.pos_encoding.embedding", "write_shape": [128, 768], "chunk_shape": [128, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.indexer.Dense_0.kernel", "write_shape": [96, 1], "chunk_shape": [96, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.indexer.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.indexer.Dense_1.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.indexer.Dense_2.bias", "write_shape": [48], "chunk_shape": [48], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.indexer.Dense_2.kernel", "write_shape": [96, 384], "chunk_shape": [96, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.indexer.Dense_3.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.indexer.Dense_3.kernel", "write_shape": [48, 1], "chunk_shape": [48, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.indexer.Dense_4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.indexer.Dense_4.kernel", "write_shape": [48, 1], "chunk_shape": [48, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.retrieval_integrator.layers_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.retrieval_integrator.layers_0.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.retrieval_integrator.layers_2.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.retrieval_integrator.layers_2.kernel", "write_shape": [96, 768], "chunk_shape": [96, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.retrieval_integrator.layers_3.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.retrieval_integrator.layers_3.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.2.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "rng", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "pool_m", "write_shape": [32768, 768], "chunk_shape": [32768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "pool_v", "write_shape": [32768, 768], "chunk_shape": [32768, 768], "ext_metadata": null}}]}

default/d/40b6413d8d045553235110cf8a3113dc

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

default/ocdbt.process_0/d/2655f27744aa28bc57a54732ca8aa17f

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

default/ocdbt.process_0/d/309bfd1f96632d6760dd55dea979babf

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

default/ocdbt.process_0/d/37b57931b1bb0df657d81dd245946279

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

default/ocdbt.process_0/d/7d3a8dd28172f4fc4fe186eaa73f2843

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

default/ocdbt.process_0/d/cca1e2cb6509e2bae33156603f3ff2de

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

infer.py

@@ -0,0 +1,840 @@
+#!/usr/bin/env python3
+"""
+DPSNR Inference — Fully self-contained single-file GPU inference for the Large model.
+
+This file contains the ENTIRE model architecture, checkpoint loading, and generation
+logic. It has ZERO dependencies on the dpsn_r_jax package.
+
+Usage:
+    source .venv/bin/activate
+
+    # Single prompt
+    python infer.py --prompt "Once upon a time"
+
+    # Interactive mode (default)
+    python infer.py
+
+    # Adjust generation parameters
+    python infer.py --prompt "The future of AI" --max_tokens 200 --temp 0.8 --top_k 50
+"""
+
+import os
+import sys
+import time
+import argparse
+from dataclasses import dataclass, field
+from collections import namedtuple
+from typing import Any, Callable, Optional
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+import jax
+import jax.numpy as jnp
+import flax.linen as nn
+from jax import lax
+from flax.training import train_state
+from flax import struct, traverse_util
+import optax
+import orbax.checkpoint
+from functools import partial
+from transformers import AutoTokenizer
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  DEVICE
+# ═══════════════════════════════════════════════════════════════════════════════
+DEVICE = jax.devices()[0]
+PLATFORM = DEVICE.platform
+print(f"[Device] {DEVICE}  (platform: {PLATFORM})")
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  CONFIG — Large model, hardcoded
+# ═══════════════════════════════════════════════════════════════════════════════
+TOKENIZER_NAME = "EleutherAI/gpt-neo-125M"
+CHECKPOINT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "checkpoints_dir")
+
+
+@dataclass
+class PoolConfig:
+    total_vectors: int
+    hidden_dim: int
+
+
+@dataclass
+class DPSNRConfig:
+    vocab_size: int = 50257
+    controller_hidden_dim: int = 768
+    controller_num_layers: int = 12
+    controller_num_heads: int = 12
+    controller_ff_multiplier: float = 2.0
+    max_seq_len: int = 1024
+    dropout: float = 0.0
+    pool_total_vectors: int = 262144
+    pool_hidden_dim: int = 768
+    librarian_hidden_dim: int = 32
+    max_reasoning_loops: int = 6
+    min_reasoning_loops: int = 1
+    halt_threshold: float = 0.99
+    min_k: int = 4
+    max_k: int = 32
+    num_clusters_to_search: int = 4
+    pad_token_id: int = 0
+    learning_rate: float = 3e-4
+    gradient_checkpointing: bool = False
+    use_bf16: bool = False
+    num_indexer_heads: int = 1
+    sigma_min: float = 0.01
+    sigma_max: float = 5.0
+    use_2d_pool: bool = False
+    pool_grid_rows: int = 512
+    pool_grid_cols: int = 512
+    sigma_anneal_steps: int = 0
+    sigma_target: float = 0.05
+    precision_loss_weight: float = 0.0
+    # Fields needed by create_train_state but unused for inference
+    streaming: bool = True
+    hf_dataset_name: Optional[str] = None
+    hf_tokenizer_name: Optional[str] = None
+    max_steps: Optional[int] = None
+    generation_steps: Optional[int] = None
+    generation_max_tokens: int = 20
+    generation_prompts: Optional[list] = None
+    num_workers: int = 4
+    loss_chunk_size: int = 0
+    finetune: Optional[Any] = None
+
+
+CONFIG = DPSNRConfig()
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  MODEL LAYERS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class FlashCausalSelfAttention(nn.Module):
+    hidden_dim: int
+    num_heads: int
+    dropout_rate: float = 0.0
+
+    @nn.compact
+    def __call__(self, x, mask=None, deterministic=True):
+        head_dim = self.hidden_dim // self.num_heads
+
+        qkv = nn.Dense(3 * self.hidden_dim, use_bias=False)(x)
+        q, k, v = jnp.split(qkv, 3, axis=-1)
+
+        q = q.reshape(x.shape[0], x.shape[1], self.num_heads, head_dim)
+        k = k.reshape(x.shape[0], x.shape[1], self.num_heads, head_dim)
+        v = v.reshape(x.shape[0], x.shape[1], self.num_heads, head_dim)
+
+        dropout_rng = (
+            self.make_rng("dropout")
+            if not deterministic and self.dropout_rate > 0
+            else None
+        )
+
+        y = nn.dot_product_attention(
+            q, k, v,
+            bias=mask,
+            dropout_rate=self.dropout_rate,
+            deterministic=deterministic,
+            dropout_rng=dropout_rng,
+        )
+
+        y = y.reshape(x.shape[0], x.shape[1], self.hidden_dim)
+        y = nn.Dense(self.hidden_dim, use_bias=False)(y)
+
+        if not deterministic:
+            y = nn.Dropout(self.dropout_rate)(y, deterministic=deterministic)
+
+        return y
+
+
+class TinyFFN(nn.Module):
+    hidden_dim: int
+    ff_dim: int
+    dropout_rate: float = 0.0
+
+    @nn.compact
+    def __call__(self, x, deterministic=True):
+        x = nn.Dense(self.ff_dim)(x)
+        x = nn.gelu(x)
+        if not deterministic:
+            x = nn.Dropout(self.dropout_rate)(x, deterministic=deterministic)
+        x = nn.Dense(self.hidden_dim)(x)
+        if not deterministic:
+            x = nn.Dropout(self.dropout_rate)(x, deterministic=deterministic)
+        return x
+
+
+class TinyTransformerLayer(nn.Module):
+    hidden_dim: int
+    num_heads: int
+    ff_dim: int
+    dropout_rate: float = 0.0
+
+    @nn.compact
+    def __call__(self, x, mask=None, deterministic=True):
+        norm1 = nn.LayerNorm()(x)
+        attn_out = FlashCausalSelfAttention(
+            self.hidden_dim, self.num_heads, self.dropout_rate
+        )(norm1, mask=mask, deterministic=deterministic)
+        x = x + attn_out
+
+        norm2 = nn.LayerNorm()(x)
+        ffn_out = TinyFFN(self.hidden_dim, self.ff_dim, self.dropout_rate)(
+            norm2, deterministic=deterministic
+        )
+        x = x + ffn_out
+        return x
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  CONTROLLER
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class TinyController(nn.Module):
+    config: DPSNRConfig
+
+    def setup(self):
+        self.embedding = nn.Embed(
+            self.config.vocab_size, self.config.controller_hidden_dim
+        )
+        self.pos_encoding = nn.Embed(
+            self.config.max_seq_len, self.config.controller_hidden_dim
+        )
+
+        ff_dim = int(
+            self.config.controller_hidden_dim * self.config.controller_ff_multiplier
+        )
+        layer_cls = TinyTransformerLayer
+        if self.config.gradient_checkpointing:
+            layer_cls = nn.remat(TinyTransformerLayer, static_argnums=(3,))
+
+        self.layers = [
+            layer_cls(
+                self.config.controller_hidden_dim,
+                self.config.controller_num_heads,
+                ff_dim,
+                self.config.dropout,
+            )
+            for _ in range(self.config.controller_num_layers)
+        ]
+
+        self.final_norm = nn.LayerNorm()
+        self.lm_head = nn.Dense(self.config.vocab_size, use_bias=False)
+
+    def __call__(self, input_ids, deterministic=True):
+        return self.encode(input_ids, deterministic)
+
+    def encode(self, input_ids, deterministic=True):
+        B, T = input_ids.shape
+        embed = self.embedding(input_ids)
+        pos_ids = jnp.arange(T)[None, :]
+        pos_embed = self.pos_encoding(pos_ids)
+        x = embed + pos_embed
+
+        mask = nn.make_causal_mask(input_ids)
+        mask = jnp.where(mask, 0, -1e4)
+
+        for layer in self.layers:
+            x = layer(x, mask, deterministic)
+
+        return x
+
+    def decode(self, hidden):
+        x = self.final_norm(hidden)
+        logits = self.lm_head(x)
+        return logits
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  MEMORY — Learned Indexer + 1D/2D Pool
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class LearnedIndexer(nn.Module):
+    hidden_dim: int
+    num_heads: int = 1
+    sigma_min: float = 0.01
+    sigma_max: float = 5.0
+
+    @nn.compact
+    def __call__(self, hidden_states, sigma_max_scale: float = 1.0):
+        attn_logits = nn.Dense(1, use_bias=False)(hidden_states)
+        attn_weights = jax.nn.softmax(attn_logits, axis=1)
+        pooled = jnp.sum(attn_weights * hidden_states, axis=1)
+
+        x = nn.Dense(self.hidden_dim)(pooled)
+        x = nn.gelu(x)
+        x = nn.Dense(self.hidden_dim // 2)(x)
+        x = nn.gelu(x)
+
+        mu_raw = nn.Dense(self.num_heads)(x)
+        sigma_raw = nn.Dense(self.num_heads)(x)
+
+        mu = jax.nn.sigmoid(mu_raw)
+
+        effective_sigma_max = self.sigma_max * sigma_max_scale
+        sigma = (
+            self.sigma_min
+            + (effective_sigma_max - self.sigma_min) * jax.nn.sigmoid(sigma_raw)
+        )
+
+        return mu, sigma
+
+
+class CoordinateMassivePool(nn.Module):
+    config: PoolConfig
+    window_size: int
+
+    def setup(self):
+        self.params_storage = self.param(
+            "params_storage",
+            nn.initializers.normal(),
+            (self.config.total_vectors, self.config.hidden_dim),
+        )
+
+    def __call__(self, mu, sigma):
+        B = mu.shape[0]
+        Total = self.config.total_vectors
+        D = self.config.hidden_dim
+        W = self.window_size
+
+        center_idx = mu * (Total - 1)
+        start_indices = jnp.clip(center_idx - W // 2, 0, Total - W).astype(jnp.int32)
+
+        def slice_fn(start):
+            return lax.dynamic_slice(self.params_storage, (start, 0), (W, D))
+
+        selected = jax.vmap(slice_fn)(start_indices)
+        relative_indices = jnp.arange(W)[None, :] + start_indices[:, None]
+        distances = relative_indices - center_idx[:, None]
+        weights = jnp.exp(-(distances**2) / (2 * (sigma[:, None] + 1e-6) ** 2)) + 1e-6
+        weights = weights / jnp.sum(weights, axis=-1, keepdims=True)
+        aggregated = jnp.einsum("bw,bwd->bd", weights, selected)
+
+        return aggregated, start_indices
+
+
+class CoordinateMassivePool2D(nn.Module):
+    rows: int
+    cols: int
+    hidden_dim: int
+    window_size: int
+
+    def setup(self):
+        self.params_storage = self.param(
+            "params_storage",
+            nn.initializers.normal(),
+            (self.rows, self.cols, self.hidden_dim),
+        )
+
+    def __call__(self, mu_row, mu_col, sigma):
+        B = mu_row.shape[0]
+        R = self.rows
+        C = self.cols
+        D = self.hidden_dim
+        W = self.window_size
+
+        r_center = mu_row * (R - 1)
+        r_start = jnp.clip(r_center - W // 2, 0, R - W).astype(jnp.int32)
+        c_center = mu_col * (C - 1)
+        c_start = jnp.clip(c_center - W // 2, 0, C - W).astype(jnp.int32)
+
+        def fetch_window(r_s, c_s):
+            return lax.dynamic_slice(self.params_storage, (r_s, c_s, 0), (W, W, D))
+
+        windows = jax.vmap(fetch_window)(r_start, c_start)
+
+        r_idx = jnp.arange(W)[None, :] + r_start[:, None]
+        c_idx = jnp.arange(W)[None, :] + c_start[:, None]
+        r_dist = r_idx - r_center[:, None]
+        c_dist = c_idx - c_center[:, None]
+
+        sigma_sq = (sigma + 1e-6) ** 2
+        r_w = jnp.exp(-r_dist ** 2 / (2 * sigma_sq[:, None]))
+        c_w = jnp.exp(-c_dist ** 2 / (2 * sigma_sq[:, None]))
+
+        w_2d = jnp.einsum("bi,bj->bij", r_w, c_w) + 1e-6
+        w_2d = w_2d / jnp.sum(w_2d, axis=(-2, -1), keepdims=True)
+
+        aggregated = jnp.einsum("bij,bijd->bd", w_2d, windows)
+        flat_start = r_start * C + c_start
+
+        return aggregated, flat_start
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  REASONING — Adaptive Compute Controller
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class AdaptiveComputeController(nn.Module):
+    hidden_dim: int
+    max_loops: int = 8
+    halt_threshold: float = 0.99
+
+    def setup(self):
+        self.halt_net = nn.Sequential(
+            [nn.Dense(self.hidden_dim // 4), nn.gelu, nn.Dense(1), nn.sigmoid]
+        )
+        self.state_gate = nn.Sequential([nn.Dense(self.hidden_dim), nn.sigmoid])
+        self.state_transform = nn.Dense(self.hidden_dim)
+        self.state_norm = nn.LayerNorm()
+        self.loop_embed = nn.Embed(32, self.hidden_dim)
+
+    def __call__(self, state_hidden, step_output, loop_count, current_halt_prob, halted_mask):
+        loop_idx = jnp.array([loop_count], dtype=jnp.int32)
+        emb = self.loop_embed(loop_idx)
+        step_output = step_output + emb
+
+        combined = jnp.concatenate([step_output, state_hidden], axis=-1)
+        g = self.state_gate(combined)
+
+        candidate_state = g * self.state_transform(step_output) + (1 - g) * state_hidden
+        candidate_state = self.state_norm(candidate_state)
+
+        hp = self.halt_net(candidate_state)
+
+        still_running_mask = 1.0 - halted_mask
+        new_halt_prob = current_halt_prob + hp * still_running_mask
+
+        is_halted_now = (new_halt_prob >= self.halt_threshold).astype(jnp.float32)
+        final_halted_mask = jnp.maximum(halted_mask, is_halted_now)
+
+        return candidate_state, new_halt_prob, final_halted_mask
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  DPSNR — Full model
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class DPSNR(nn.Module):
+    config: DPSNRConfig
+
+    def setup(self):
+        self.controller = TinyController(self.config)
+        self.indexer = LearnedIndexer(
+            self.config.controller_hidden_dim,
+            num_heads=self.config.num_indexer_heads,
+            sigma_min=self.config.sigma_min,
+            sigma_max=self.config.sigma_max,
+        )
+
+        if self.config.use_2d_pool:
+            axis_window = max(2, int(self.config.max_k ** 0.5))
+            self.pool = CoordinateMassivePool2D(
+                rows=self.config.pool_grid_rows,
+                cols=self.config.pool_grid_cols,
+                hidden_dim=self.config.controller_hidden_dim,
+                window_size=axis_window,
+            )
+        else:
+            self.pool = CoordinateMassivePool(
+                PoolConfig(
+                    self.config.pool_total_vectors,
+                    self.config.controller_hidden_dim,
+                ),
+                window_size=self.config.max_k,
+            )
+
+        self.acc = AdaptiveComputeController(
+            self.config.controller_hidden_dim,
+            self.config.max_reasoning_loops,
+            self.config.halt_threshold,
+        )
+        self.retrieval_integrator = nn.Sequential(
+            [
+                nn.Dense(self.config.controller_hidden_dim),
+                nn.gelu,
+                nn.Dense(self.config.controller_hidden_dim),
+                nn.LayerNorm(),
+            ]
+        )
+
+    def __call__(self, input_ids, deterministic=True, sigma_max_scale: float = 1.0):
+        state_hidden, all_indices, mean_sigma = self._encode_hidden(
+            input_ids, deterministic, sigma_max_scale
+        )
+        logits = self.controller.decode(state_hidden)
+        return logits, (self.config.max_reasoning_loops, all_indices, mean_sigma)
+
+    def encode_to_hidden(self, input_ids, deterministic=True, sigma_max_scale: float = 1.0):
+        state_hidden, all_indices, mean_sigma = self._encode_hidden(
+            input_ids, deterministic, sigma_max_scale
+        )
+        return state_hidden, (self.config.max_reasoning_loops, all_indices, mean_sigma)
+
+    def _encode_hidden(self, input_ids, deterministic=True, sigma_max_scale: float = 1.0):
+        hidden = self.controller(input_ids, deterministic)
+        state_hidden = hidden
+        B, T, D = hidden.shape
+
+        halt_prob = jnp.zeros((B, T, 1), dtype=hidden.dtype)
+        halted_mask = jnp.zeros((B, T, 1), dtype=hidden.dtype)
+
+        # Warm-up calls: force Flax to trace all sub-modules before scan
+        _mu, _sigma = self.indexer(
+            jnp.zeros((B, T, D)), sigma_max_scale=sigma_max_scale
+        )
+        if self.config.use_2d_pool:
+            H = self.config.num_indexer_heads
+            h_per_dim = max(1, H // 2)
+            _ = self.pool(jnp.zeros((B,)), jnp.zeros((B,)), jnp.zeros((B,)))
+        else:
+            _ = self.pool(jnp.zeros((B,)), jnp.zeros((B,)))
+        _ = self.retrieval_integrator(
+            jnp.zeros((B, T, D + self.config.controller_hidden_dim))
+        )
+        _ = self.acc(state_hidden, state_hidden, 0, halt_prob, halted_mask)
+
+        use_2d = self.config.use_2d_pool
+        H = self.config.num_indexer_heads
+
+        def reasoning_step(carry, i):
+            s_hidden, h_prob, h_mask = carry
+            prev_s_hidden = s_hidden
+
+            mu, sigma = self.indexer(s_hidden, sigma_max_scale=sigma_max_scale)
+
+            all_retrieved = []
+            all_start_indices = []
+
+            if use_2d:
+                heads_per_dim = max(1, H // 2)
+                for h in range(heads_per_dim):
+                    h_row = h
+                    h_col = min(h + heads_per_dim, H - 1)
+                    sigma_h = (sigma[:, h_row] + sigma[:, h_col]) / 2.0
+                    retrieved_h, start_idx_h = self.pool(
+                        mu[:, h_row], mu[:, h_col], sigma_h
+                    )
+                    all_retrieved.append(retrieved_h)
+                    all_start_indices.append(start_idx_h)
+            else:
+                for h in range(H):
+                    retrieved_h, start_idx_h = self.pool(mu[:, h], sigma[:, h])
+                    all_retrieved.append(retrieved_h)
+                    all_start_indices.append(start_idx_h)
+
+            retrieved = jnp.mean(jnp.stack(all_retrieved, axis=1), axis=1)
+            start_indices = jnp.concatenate(all_start_indices, axis=0)
+            mean_sigma_step = jnp.mean(sigma)
+
+            retrieved_expanded = jnp.expand_dims(retrieved, 1).repeat(T, axis=1)
+            combined = jnp.concatenate([s_hidden, retrieved_expanded], axis=-1)
+            integrated = self.retrieval_integrator(combined)
+
+            new_s_hidden, h_prob, new_h_mask = self.acc(
+                s_hidden, s_hidden + integrated, i, h_prob, h_mask,
+            )
+
+            update_mask = 1.0 - h_mask
+            s_hidden = update_mask * new_s_hidden + h_mask * prev_s_hidden
+
+            carry_dtype = prev_s_hidden.dtype
+            s_hidden = s_hidden.astype(carry_dtype)
+            h_prob = h_prob.astype(carry_dtype)
+            new_h_mask = new_h_mask.astype(carry_dtype)
+
+            return (s_hidden, h_prob, new_h_mask), (start_indices, mean_sigma_step)
+
+        _scan_fn = reasoning_step
+        if self.config.gradient_checkpointing:
+            _scan_fn = jax.checkpoint(reasoning_step)
+
+        init_carry = (state_hidden, halt_prob, halted_mask)
+        (state_hidden, halt_prob, halted_mask), (all_indices, sigma_per_loop) = (
+            jax.lax.scan(
+                _scan_fn,
+                init_carry,
+                jnp.arange(self.config.max_reasoning_loops),
+            )
+        )
+
+        all_indices = jnp.transpose(all_indices, (1, 0))
+        mean_sigma = jnp.mean(sigma_per_loop)
+
+        return state_hidden, all_indices, mean_sigma
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  TRAIN STATE — Minimal, just enough to restore the checkpoint pytree
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class TrainState(train_state.TrainState):
+    rng: Any
+    pool_m: jnp.ndarray
+    pool_v: jnp.ndarray
+    window_size: int = struct.field(pytree_node=False)
+    learning_rate_fn: Callable[[int], float] = struct.field(pytree_node=False)
+    sigma_anneal_fn: Callable[[int], float] = struct.field(pytree_node=False)
+
+
+def _create_dummy_state(rng, config):
+    """Create a dummy TrainState with the correct pytree structure for checkpoint restore."""
+    model = DPSNR(config)
+    dummy_input = jnp.ones((1, config.max_seq_len), dtype=jnp.int32)
+    variables = model.init(rng, dummy_input)
+    params = variables["params"]
+
+    flat_params = traverse_util.flatten_dict(params)
+    pool_key = ("pool", "params_storage")
+    pool_params = flat_params[pool_key]
+
+    dense_flat_params = {k: v for k, v in flat_params.items() if k != pool_key}
+    dense_params = traverse_util.unflatten_dict(dense_flat_params)
+
+    learning_rate_fn = lambda step: config.learning_rate
+
+    tx = optax.chain(
+        optax.clip_by_global_norm(1.0),
+        optax.adamw(learning_rate=learning_rate_fn),
+    )
+    opt_state = tx.init(dense_params)
+
+    pool_m = jnp.zeros_like(pool_params)
+    pool_v = jnp.zeros_like(pool_params)
+
+    sigma_anneal_fn = lambda step: 1.0
+
+    return TrainState(
+        step=jnp.array(0, dtype=jnp.int32),
+        apply_fn=model.apply,
+        params=params,
+        tx=tx,
+        opt_state=opt_state,
+        rng=rng,
+        pool_m=pool_m,
+        pool_v=pool_v,
+        window_size=config.max_k,
+        learning_rate_fn=learning_rate_fn,
+        sigma_anneal_fn=sigma_anneal_fn,
+    )
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  INFERENCE CONTAINER
+# ═══════════════════════════════════════════════════════════════════════════════
+InferenceModel = namedtuple("InferenceModel", ["apply_fn", "params", "step"])
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  TOKENIZER
+# ═══════════════════════════════════════════════════════════════════════════════
+def load_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    return tokenizer
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  CHECKPOINT LOADING
+# ═══════════════════════════════════════════════════════════════════════════════
+def load_checkpoint():
+    """Load trained weights from checkpoint. Returns only params + apply_fn."""
+    rng = jax.random.PRNGKey(0)
+
+    cpu = jax.devices("cpu")[0]
+    print("[Init] Creating model skeleton on CPU...")
+    with jax.default_device(cpu):
+        dummy_state = _create_dummy_state(rng, CONFIG)
+    dummy_state = jax.device_put(dummy_state, cpu)
+
+    abs_ckpt = os.path.abspath(CHECKPOINT_DIR)
+    checkpointer = orbax.checkpoint.PyTreeCheckpointer()
+    restore_args = orbax.checkpoint.checkpoint_utils.construct_restore_args(dummy_state)
+
+    mgr = orbax.checkpoint.CheckpointManager(abs_ckpt, checkpointer)
+    latest_step = mgr.latest_step()
+
+    if latest_step is not None:
+        print(f"[Checkpoint] Restoring step {latest_step} from {abs_ckpt}")
+        state = mgr.restore(
+            latest_step,
+            items=dummy_state,
+            restore_kwargs={"restore_args": restore_args},
+        )
+    else:
+        target = None
+        for sub in ("default", ""):
+            p = os.path.join(abs_ckpt, sub) if sub else abs_ckpt
+            if os.path.exists(os.path.join(p, "_METADATA")):
+                target = p
+                break
+        if target is None:
+            raise FileNotFoundError(f"No valid checkpoint found in {abs_ckpt}")
+        print(f"[Checkpoint] Restoring directly from {target}")
+        state = checkpointer.restore(target, item=dummy_state, restore_args=restore_args)
+
+    step = int(state.step)
+    apply_fn = state.apply_fn
+    params = state.params
+
+    del dummy_state, state
+
+    if PLATFORM != "cpu":
+        print(f"[Device] Moving model params to {DEVICE}...")
+        params = jax.device_put(params, DEVICE)
+
+    print(f"[Checkpoint] Loaded at training step {step}")
+    return InferenceModel(apply_fn=apply_fn, params=params, step=step)
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  JIT FORWARD PASS
+# ═══════════════════════════════════════════════════════════════════════════════
+@partial(jax.jit, static_argnums=(0,))
+def _forward(apply_fn, params, input_ids):
+    logits, _ = apply_fn({"params": params}, input_ids, deterministic=True)
+    return logits
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  TEXT GENERATION
+# ═══════════════════════════════════════════════════════════════════════════════
+def generate(
+    model: InferenceModel,
+    prompt: str,
+    tokenizer,
+    rng,
+    max_tokens: int = 100,
+    temperature: float = 0.7,
+    top_k: int = 40,
+    repetition_penalty: float = 1.2,
+):
+    """Autoregressive generation with fixed-size buffers (no XLA recompilation)."""
+    input_ids = tokenizer.encode(prompt, return_tensors="np")
+    eos_id = tokenizer.eos_token_id
+    prompt_len = input_ids.shape[1]
+    max_seq = CONFIG.max_seq_len
+
+    if prompt_len > max_seq:
+        input_ids = input_ids[:, :max_seq]
+        prompt_len = max_seq
+
+    buf = jnp.zeros((1, max_seq), dtype=jnp.int32)
+    buf = buf.at[:, :prompt_len].set(input_ids)
+
+    gen_buf = jnp.zeros((max_tokens,), dtype=jnp.int32)
+    n_gen = 0
+
+    for step in range(max_tokens):
+        pos = prompt_len + step
+        if pos >= max_seq:
+            break
+
+        logits = _forward(model.apply_fn, model.params, buf)
+        next_logits = logits[0, pos - 1, :]
+
+        # Repetition penalty
+        if n_gen > 0:
+            prev = gen_buf[:n_gen]
+            vocab = next_logits.shape[-1]
+            mask = jnp.zeros(vocab, dtype=jnp.bool_)
+            mask = mask.at[prev].set(True)
+            penalized = jnp.where(
+                next_logits > 0,
+                next_logits / repetition_penalty,
+                next_logits * repetition_penalty,
+            )
+            next_logits = jnp.where(mask, penalized, next_logits)
+
+        # Top-k filtering
+        k = min(top_k, next_logits.shape[-1])
+        vals, _ = jax.lax.top_k(next_logits, k=k)
+        threshold = vals[-1]
+        next_logits = jnp.where(next_logits < threshold, -1e10, next_logits)
+
+        # Temperature sampling
+        rng, key = jax.random.split(rng)
+        token = jax.random.categorical(key, next_logits / max(temperature, 1e-8))
+        token_int = int(token)
+
+        buf = buf.at[0, pos].set(token_int)
+        gen_buf = gen_buf.at[n_gen].set(token_int)
+        n_gen += 1
+
+        if token_int == eos_id:
+            break
+
+    return tokenizer.decode(
+        buf[0, prompt_len : prompt_len + n_gen].tolist(),
+        skip_special_tokens=True,
+    )
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+#  MAIN
+# ═══════════════════════════════════════════════════════════════════════════════
+def main():
+    parser = argparse.ArgumentParser(description="DPSNR Large — Inference")
+    parser.add_argument("--prompt", type=str, default=None, help="Input prompt (omit for interactive mode)")
+    parser.add_argument("--max_tokens", type=int, default=100, help="Max tokens to generate (default: 100)")
+    parser.add_argument("--temp", type=float, default=0.7, help="Sampling temperature (default: 0.7)")
+    parser.add_argument("--top_k", type=int, default=40, help="Top-k sampling (default: 40)")
+    parser.add_argument("--penalty", type=float, default=1.2, help="Repetition penalty (default: 1.2)")
+    parser.add_argument("--checkpoint_dir", type=str, default=None, help="Override checkpoint path")
+    args = parser.parse_args()
+
+    if args.checkpoint_dir:
+        global CHECKPOINT_DIR
+        CHECKPOINT_DIR = args.checkpoint_dir
+
+    print("=" * 60)
+    print("  DPSNR Large — Loading Model")
+    print("=" * 60)
+    tokenizer = load_tokenizer()
+    model = load_checkpoint()
+
+    # Warmup: compile forward pass once
+    print("[Warmup] Compiling forward pass...")
+    t0 = time.time()
+    warmup_ids = jnp.zeros((1, CONFIG.max_seq_len), dtype=jnp.int32)
+    _ = _forward(model.apply_fn, model.params, warmup_ids)
+    jax.effects_barrier()
+    print(f"[Warmup] Done in {time.time() - t0:.1f}s")
+
+    rng = jax.random.PRNGKey(42)
+
+    def run(prompt: str):
+        nonlocal rng
+        rng, key = jax.random.split(rng)
+        t0 = time.time()
+        output = generate(
+            model, prompt, tokenizer, key,
+            max_tokens=args.max_tokens,
+            temperature=args.temp,
+            top_k=args.top_k,
+            repetition_penalty=args.penalty,
+        )
+        elapsed = time.time() - t0
+        print(f"\n{'─' * 50}")
+        print(f"Prompt:    {prompt}")
+        print(f"Generated: {output}")
+        print(f"Time:      {elapsed:.2f}s")
+        print(f"{'─' * 50}\n")
+
+    if args.prompt:
+        run(args.prompt)
+    else:
+        print("\n╔══════════════════════════════════════════════════╗")
+        print("║   DPSNR Interactive Inference                    ║")
+        print("║   Type 'exit' or 'quit' to stop                  ║")
+        print("╚══════════════════════════════════════════════════╝\n")
+        while True:
+            try:
+                user_input = input(">>> ")
+                if user_input.strip().lower() in ("exit", "quit"):
+                    break
+                if not user_input.strip():
+                    continue
+                run(user_input)
+            except (EOFError, KeyboardInterrupt):
+                print("\nExiting...")
+                break
+
+
+if __name__ == "__main__":
+    main()