Publish audited ActionCodec2 release files

README.md

@@ -9,32 +9,16 @@ tags:
 
 # ActionCodec2
 
-Tokenizers for **ActionCodec2: A Streaming Action Tokenizer**. This repository contains
-30 variants with one shared runtime and an optional native extension.
+- **Paper edition**: first- and second-order tokenizers used in the paper; 5 mm translation precision and 4,096 tokens per profile.
+- **Extended edition**: retrained after additional pretraining-data cleaning and filtering, with more precision and vocabulary options.
+- **Sparse version**: standard and high-dynamic first-order tokenizers, each with EEF and joint profiles and binary (`2bin`) or 21-level (`21bin`) grippers.
 
-## Editions
-
-- **Paper edition**: `paper-1st` and `paper-2nd` are the first- and second-order tokenizers
-  used in the paper. Both use **5 mm translation precision** and a **4,096-token budget
-  per profile**, with joint and EEF profiles.
-- **Extended edition**: tokenizers retrained after additional cleaning and filtering of
-  the pretraining data. They offer first and second order, **1/5/10 mm translation
-  precision**, and **1k–16k token budgets per profile** for a wider range of precision
-  and vocabulary requirements. The 1 mm tier starts at 2k; its 2k variants contain
-  EEF profiles only.
-
-Precision labels describe translation quantization settings. Rotation and joint settings
-are specified separately in each profile's configuration. Paper and Extended variants
-have separately trained vocabularies, including when their precision and budget match.
-The six released VLA policies use the Paper edition and must keep their original variant.
+All 34 variants use one shared runtime. Precision labels describe translation quantization; rotation and joint limits are in each profile's configuration.
 
 ## Tokenizer usage
 
-Select a variant explicitly to load its configuration and fitted assets together:
-
 ```bash
-export HF_ENDPOINT=https://hf-mirror.com
-pip install numpy scipy torch "transformers>=4.57,<6" huggingface-hub pyyaml
+pip install numpy scipy torch "transformers>=5.2,<6" huggingface-hub pyyaml
 ```
 
 ```python
@@ -42,50 +26,46 @@ import numpy as np
 from transformers import AutoProcessor
 
 codec = AutoProcessor.from_pretrained(
-    "ZibinDong/ActionCodec2", variant="paper-1st", trust_remote_code=True,
-    revision="v0.1.1",
+    "ZibinDong/ActionCodec2", variant="paper-1st",
+    action_space="single_eef_delta", trust_remote_code=True,
 )
-actions = np.zeros((12, 7), dtype=np.float32)
+actions = np.zeros((1, 16, 7), dtype=np.float32)
 tokens = codec.encode(actions, fps=15)
-decoded = codec.decode(tokens, fps=15)
-codec.save_pretrained("./saved-tokenizer")
-restored = AutoProcessor.from_pretrained(
-    "./saved-tokenizer", trust_remote_code=True, local_files_only=True,
-)
+restored = codec.decode(tokens, fps=15)
 ```
 
-The default action space is `single_eef_delta`: translation increments in metres,
-body-frame rotation-vector increments in radians, and an absolute binary gripper command.
-Rotation follows `R_next = R_previous @ Exp(rotvec)`; the gripper opens at values >= 0.8.
-Pass the actual recording rate as `fps`; the tokenizer resamples to 15 Hz.
-Inputs have shape `(T, D)` or `(B, T, D)`; decoded CPU float32 tensors have shape `(B, T_out, D)`.
-Resampling may change the length. Quantization is lossy.
+Use `variant="sparse-standard-2bin"` or `"sparse-high-dynamic-21bin"` for Sparse;
+Sparse decoding requires `horizon=actions.shape[-2]` in source frames.
+Each profile has its own token namespace; the budget is per profile, not the total routed vocabulary.
+
+The EEF layout is `[x, y, z, rx, ry, rz, gripper]`, in metres and radians.
+Delta rotations use `R_next = R_previous @ Exp(rotvec)`. Grippers are absolute open fractions (0 closed, 1 open).
+Inputs are `(T,D)` or `(B,T,D)` and are resampled to 15 Hz. Reconstruction is lossy.
+Absolute layouts also require `current_state`; second-order variants require
+`previous_first_order` in canonical 15 Hz increments in both encode and decode.
+Zero boundaries are valid only when the preceding motion is known to be zero.
+Use `final_first_order` to carry reconstructed second-order boundaries between chunks.
 
-Second-order variants require the preceding first-order motion at every chunk boundary:
+`codec.print_action_spaces()` lists the layouts. Dual arms are interleaved by time;
+BPE tokens never merge across arms. The historical `single_joint6_abs_continuous_gripper`
+name aliases `single_joint6_abs_no_gripper`: six joint coordinates, no separate gripper slot.
+
+## Train tokenizers
 
 ```python
-codec = AutoProcessor.from_pretrained(
-    "ZibinDong/ActionCodec2", variant="paper-2nd", trust_remote_code=True,
-    revision="v0.1.1",
-)
-# Use a zero template only when the preceding motion is known to be zero.
-boundary = codec.first_order_template(batch_size=1)
-tokens = codec.encode(actions, fps=15, previous_first_order=boundary)
-decoded = codec.decode(tokens, fps=15, previous_first_order=boundary)
-boundary = codec.final_first_order(tokens, previous_first_order=boundary, executed_steps=None)
+codec.fit(episodes, fps=15)
+codec.save_pretrained("./my-tokenizer")
 ```
 
-`executed_steps` counts reconstructed steps at the tokenizer's internal **15 Hz**, even
-when `encode` or `decode` uses a different `fps`; `None` carries the full chunk boundary.
-`previous_first_order` uses the same 15 Hz convention. The VLA adapter handles conversion
-from the last raw executed action at the policy's control frequency.
-Use `codec.for_action_space("single_joint7_delta")` or another registered layout to change
-the physical input convention. Absolute layouts require `current_state` for reconstruction.
-The 1mm/2k variants contain EEF profiles only.
+Cleaning is off by default, including VLA training and inference. Only tokenizer fitting
+accepts `apply_preprocessing=True`, with a recorded or explicitly supplied `preprocessing`
+configuration. It acts on continuous canonical trajectories before quantization.
+Split recording discontinuities first; second-order cleaning needs `bootstrap_context=True`.
+Saved cleaning settings do not enable filtering during encode or decode.
 
 ## Available variants
 
-| Variant | Edition | Order | Translation precision | Token budget per profile | Profiles |
+| Variant | Edition | Order | Translation precision | Budget / profile | Profiles |
 | --- | --- | --- | --- | --- | --- |
 | `paper-1st` | Paper | 1 | 5mm | 4096 | joint, eef |
 | `paper-2nd` | Paper | 2 | 5mm | 4096 | joint, eef |
@@ -117,52 +97,22 @@ The 1mm/2k variants contain EEF profiles only.
 | `extended-2nd-10mm-4k` | Extended | 2 | 10mm | 4096 | joint, eef |
 | `extended-2nd-10mm-8k` | Extended | 2 | 10mm | 8192 | joint, eef |
 | `extended-2nd-10mm-16k` | Extended | 2 | 10mm | 16384 | joint, eef |
+| `sparse-standard-2bin` | Sparse | 1 | 5mm | 4096 | eef, joint |
+| `sparse-standard-21bin` | Sparse | 1 | 5mm | 4096 | eef, joint |
+| `sparse-high-dynamic-2bin` | Sparse | 1 | 5mm | 4096 | eef, joint |
+| `sparse-high-dynamic-21bin` | Sparse | 1 | 5mm | 4096 | eef, joint |
 
-Joint and EEF token IDs occupy disjoint namespaces. Budgets apply separately to each
-profile. Quantization is lossy; the precision label alone does not determine policy accuracy.
-Pin `revision` to a commit for reproducibility; `variants.json` records each variant's
-asset hashes and its previous repository ID. Loading downloads shared files plus the
-selected variant, and verifies its assets. `subfolder` is not a variant selector.
-
-Saved processors contain only the selected assets and the complete portable runtime.
-They load offline and preserve `artifact_manifest.json` with their asset hashes and origin.
-Select a different variant by loading the catalog again.
-
-## Native acceleration
-
-Build the optional C++17 kernels once for all variants:
-
-```bash
-hf download ZibinDong/ActionCodec2 --revision v0.1.1 --include "setup.py" "pyproject.toml" "LICENSE" "actioncodec2_kernels/*" --local-dir ./actioncodec2-kernels
-pip install ./actioncodec2-kernels
-```
-
-`ACTIONCODEC2_NUM_THREADS` controls the per-process thread budget.
-`ACTIONCODEC2_NO_NATIVE=1` selects the Python reference implementation.
-
-## VLA inference package
-
-The versioned source archive `actioncodec2-0.1.1.tar.gz` contains policy inference code,
-native kernels, training recipes, and evaluation scripts. Use **Python 3.12 or later**
-and a C++17 compiler. Inference has been checked with Python 3.12, PyTorch 2.7.1,
-and Transformers 5.14.1.
-
-```bash
-pip install "actioncodec2[cuda] @ https://hf-mirror.com/ZibinDong/ActionCodec2/resolve/v0.1.1/actioncodec2-0.1.1.tar.gz"
-```
+Pin `revision` for reproducibility. `variants.json` records asset hashes; a selected
+`save_pretrained()` directory contains its assets and the shared portable runtime.
+The optional C++17 kernels can be built with `pip install ./downloaded-tokenizer`.
 
-The `cuda` extra adds Liger and Flash Linear Attention. Each policy bundles its exact
-AC2 variant and text tokenizer; once downloaded, inference works offline without
-another model or tokenizer repository. Each policy's README describes its inputs.
+## VLA policies
 
-| Training data / evaluation target | First-order policy | Second-order policy |
-| --- | --- | --- |
-| LIBERO | [First order](https://hf-mirror.com/ZibinDong/G05_ActionCodec2-1st_LIBERO) | [Second order](https://hf-mirror.com/ZibinDong/G05_ActionCodec2-2nd_LIBERO) |
-| Bridge / SimplerEnv (WidowX) | [First order](https://hf-mirror.com/ZibinDong/G05_ActionCodec2-1st_SimplerEnv) | [Second order](https://hf-mirror.com/ZibinDong/G05_ActionCodec2-2nd_SimplerEnv) |
-| DROID | [First order](https://hf-mirror.com/ZibinDong/G05_ActionCodec2-1st_DROID) | [Second order](https://hf-mirror.com/ZibinDong/G05_ActionCodec2-2nd_DROID) |
+Installation, paper-policy evaluation and VLA training: [source repository](https://github.com/ZibinDong/actioncodec2).
+The existing six AC2 policies retain their single-arm tokenizer behavior. The released Sparse
+profiles have a 4,096-token budget; the current G0.5 Sparse adapter additionally needs an EOS
+position and therefore requires a vocabulary of at most 4,095 tokens when training a new VLA.
 
 ## License
 
-The tokenizer runtime and fitted assets use [Apache-2.0](LICENSE). The source archive
-also includes G0.5 components under `LICENSE-G0.5`. G0.5-derived policy weights use the
-G0.5 Community License; see the individual policy cards and the archive's license files.
+Tokenizer code and assets: Apache-2.0. G0.5 policy code and weights use their separate community license.

actioncodec2-0.1.1.tar.gz → actioncodec2-0.3.0.tar.gz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dac31a13dfbe46534186f8b51db19fceab9febcfe54fd880f029b5fd9ce5915d
-size 302806
+oid sha256:fdb76fb260435f45fba5a6572122eee3f3ceacb11480091f781e645ab60c0bec
+size 503117

actioncodec2_kernels/__init__.py

@@ -1 +0,0 @@
-"""Portable ActionCodec2 CPU kernels."""

actioncodec2_kernels/bpe.cpp

@@ -0,0 +1,428 @@
+// Native BPE trainer -- the FAST-style ablation baseline.
+//
+// Set-BPE treats one control step as an unordered set of primitive cells and
+// merges spatially before merging across time.  Plain BPE instead serializes
+// the quantized ``[T, D]`` cell matrix row-major into one symbol sequence and
+// runs textbook byte-pair encoding on it, which is what FAST does after its
+// DCT stage.  Both trainers consume the identical quantized corpus,
+// so a paired run isolates the tokenizer structure and nothing else.
+//
+// Correctness is identical to the Python reference trainer in
+// ``actioncodec2/bpe_training.py``, including the tie-break (largest count,
+// then lexicographically smallest ordered ``(left, right)`` pair) and the
+// left-to-right application order, so the two emit the same merge list.
+//
+// Where the time goes, and what is done about it:
+//   * occurrences of every observed ordered pair are indexed once, so a merge
+//     touches only the positions that actually carry it;
+//   * the corpus is a doubly linked list over one flat symbol array, so
+//     collapsing a pair is O(1) and never rewrites the tail;
+//   * counters live in the same flat linear-probing table as the Set-BPE
+//     kernels (common.hpp), and the heap is lazy with per-merge deduplication
+//     of the keys a merge could have changed.
+
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <pybind11/stl.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#include "common.hpp"
+
+namespace py = pybind11;
+
+namespace {
+
+// BPE pairs are ordered, so the key keeps ``left`` in the high word.  The
+// heap tie-break in common.hpp prefers the smallest key, which is therefore the
+// lexicographically smallest ``(left, right)`` -- the Python reference rule.
+inline uint64_t ordered_key(int32_t left, int32_t right) {
+  return (static_cast<uint64_t>(static_cast<uint32_t>(left)) << 32) |
+         static_cast<uint64_t>(static_cast<uint32_t>(right));
+}
+
+inline int32_t key_left(uint64_t key) { return static_cast<int32_t>(key >> 32); }
+inline int32_t key_right(uint64_t key) { return static_cast<int32_t>(key & 0xffffffffull); }
+
+// Per-merge set of pair keys whose counters moved.  A merge can only touch
+// pairs incident to its two parents or its child, so the live population is
+// bounded by a small multiple of the vocabulary; clearing is a generation bump
+// rather than a table sweep.
+class StampSet {
+ public:
+  explicit StampSet(size_t capacity_hint) {
+    size_t target = 1024;
+    while (target < capacity_hint * 4) target <<= 1;
+    keys_.assign(target, 0);
+    stamps_.assign(target, 0);
+    mask_ = target - 1;
+  }
+
+  void reset() {
+    ++generation_;
+    members_.clear();
+    if (generation_ != 0) return;
+    std::fill(stamps_.begin(), stamps_.end(), 0);
+    generation_ = 1;
+  }
+
+  void insert(uint64_t key) {
+    size_t slot = mix(key) & mask_;
+    while (stamps_[slot] == generation_) {
+      if (keys_[slot] == key) return;
+      slot = (slot + 1) & mask_;
+    }
+    if ((members_.size() + 1) * 4 > (mask_ + 1) * 3) {
+      grow();
+      slot = mix(key) & mask_;
+      while (stamps_[slot] == generation_) {
+        if (keys_[slot] == key) return;
+        slot = (slot + 1) & mask_;
+      }
+    }
+    keys_[slot] = key;
+    stamps_[slot] = generation_;
+    members_.push_back(key);
+  }
+
+  const std::vector<uint64_t>& members() const { return members_; }
+
+ private:
+  static uint64_t mix(uint64_t value) {
+    value ^= value >> 30;
+    value *= 0xbf58476d1ce4e5b9ull;
+    value ^= value >> 27;
+    value *= 0x94d049bb133111ebull;
+    return value ^ (value >> 31);
+  }
+
+  void grow() {
+    const size_t buckets = (mask_ + 1) << 1;
+    keys_.assign(buckets, 0);
+    stamps_.assign(buckets, 0);
+    mask_ = buckets - 1;
+    generation_ = 1;
+    for (const uint64_t key : members_) {
+      size_t slot = mix(key) & mask_;
+      while (stamps_[slot] == generation_) slot = (slot + 1) & mask_;
+      keys_[slot] = key;
+      stamps_[slot] = generation_;
+    }
+  }
+
+  std::vector<uint64_t> keys_;
+  std::vector<uint32_t> stamps_;
+  std::vector<uint64_t> members_;
+  size_t mask_ = 0;
+  uint32_t generation_ = 1;
+};
+
+// Occurrence index.  ``slots`` maps a pair key to its 1-based row in
+// ``lists``; a value of zero means "no row yet", which is exactly the absent
+// value CounterMap already reports, so no separate presence bit is needed.
+struct OccurrenceIndex {
+  ac2::CounterMap slots{1 << 16};
+  std::vector<std::vector<int32_t>> lists;
+  std::vector<uint64_t> keys;
+  int64_t stored = 0;
+
+  int64_t row_for(uint64_t key) {
+    const int64_t existing = slots.get(key);
+    if (existing > 0) return existing - 1;
+    const int64_t row = static_cast<int64_t>(lists.size());
+    lists.emplace_back();
+    keys.push_back(key);
+    slots.add(key, row + 1);
+    return row;
+  }
+
+  void append(uint64_t key, int32_t position) {
+    lists[static_cast<size_t>(row_for(key))].push_back(position);
+    ++stored;
+  }
+};
+
+struct Corpus {
+  std::vector<int32_t> symbol;  // live token id per slot, -1 once absorbed
+  std::vector<int32_t> prev;
+  std::vector<int32_t> next;
+  int64_t actions = 0;
+  int64_t episodes = 0;
+};
+
+void ingest_chunk(Corpus& corpus, const int32_t* cells, int64_t rows, int32_t dimension,
+                  const int32_t* offsets, const int32_t* bins, const int64_t* episode_offsets,
+                  int64_t episode_count) {
+  for (int64_t episode = 0; episode < episode_count; ++episode) {
+    const int64_t start = episode_offsets[episode];
+    const int64_t stop = episode_offsets[episode + 1];
+    if (stop <= start || start < 0 || stop > rows) {
+      throw std::invalid_argument("BPE episode offsets must be increasing and in range");
+    }
+    const size_t first = corpus.symbol.size();
+    const int32_t* initial = cells + start * dimension;
+    for (int64_t row = start; row < stop; ++row) {
+      const int32_t* values = cells + row * dimension;
+      for (int32_t dim = 0; dim < dimension; ++dim) {
+        const int32_t value = values[dim];
+        if ((value >= 0) != (initial[dim] >= 0)) {
+          throw std::invalid_argument("active dimensions must remain constant within an episode");
+        }
+        if (value < 0) {
+          if (value != -1) throw std::invalid_argument("cells must contain -1 or a valid index");
+          continue;
+        }
+        if (value >= bins[dim]) throw std::invalid_argument("cell index exceeds its bin count");
+        corpus.symbol.push_back(offsets[dim] + value);
+      }
+    }
+    const size_t last = corpus.symbol.size();
+    if (last == first) throw std::invalid_argument("every control step needs an active coordinate");
+    if (last > static_cast<size_t>(INT32_MAX)) {
+      throw std::invalid_argument(
+          "BPE corpus exceeds 2^31 primitive cells; lower the fitting coverage");
+    }
+    corpus.prev.resize(last);
+    corpus.next.resize(last);
+    for (size_t index = first; index < last; ++index) {
+      corpus.prev[index] = index == first ? -1 : static_cast<int32_t>(index - 1);
+      corpus.next[index] = index + 1 == last ? -1 : static_cast<int32_t>(index + 1);
+    }
+    corpus.actions += stop - start;
+    ++corpus.episodes;
+  }
+}
+
+}  // namespace
+
+py::tuple train_bpe_native(py::iterable cell_chunks, const std::vector<int32_t>& bins,
+                            int64_t vocab_size, int64_t max_token_cells, int threads) {
+  const int32_t dimension = static_cast<int32_t>(bins.size());
+  if (dimension < 1) throw std::invalid_argument("bins must not be empty");
+  std::vector<int32_t> offsets(dimension + 1, 0);
+  for (int32_t index = 0; index < dimension; ++index) {
+    if (bins[index] < 1) throw std::invalid_argument("every dimension needs at least one bin");
+    offsets[index + 1] = offsets[index] + bins[index];
+  }
+  const int32_t atom_count = offsets[dimension];
+  if (vocab_size < atom_count) {
+    throw std::invalid_argument("vocab_size is below the number of primitive atoms");
+  }
+  if (max_token_cells < 1) throw std::invalid_argument("max_token_cells must be positive");
+
+  // Greedy merge updates are serial. Preserve the requested option for the
+  // shared trainer API without changing the calling thread's OpenMP settings.
+  const int thread_count = 1;
+
+  Corpus corpus;
+  for (py::handle raw_chunk : cell_chunks) {
+    py::tuple chunk = py::cast<py::tuple>(raw_chunk);
+    if (chunk.size() != 2) throw std::invalid_argument("each chunk must be (cells, offsets)");
+    auto cells = py::cast<py::array_t<int32_t, py::array::c_style | py::array::forcecast>>(chunk[0]);
+    auto bounds =
+        py::cast<py::array_t<int64_t, py::array::c_style | py::array::forcecast>>(chunk[1]);
+    auto cells_view = cells.unchecked<2>();
+    auto bounds_view = bounds.unchecked<1>();
+    if (cells_view.shape(1) != dimension) {
+      throw std::invalid_argument("cells must have one column per bin count");
+    }
+    if (bounds_view.shape(0) < 2) throw std::invalid_argument("offsets must bound one episode");
+    {
+      py::gil_scoped_release release;
+      ingest_chunk(corpus, cells_view.data(0, 0), cells_view.shape(0), dimension, offsets.data(),
+                   bins.data(), bounds_view.data(0), bounds_view.shape(0) - 1);
+    }
+  }
+  if (corpus.symbol.empty()) throw std::invalid_argument("BPE training needs one episode");
+
+  const int64_t initial_length = static_cast<int64_t>(corpus.symbol.size());
+  int64_t length = initial_length;
+  std::vector<int32_t>& symbol = corpus.symbol;
+  std::vector<int32_t>& prev = corpus.prev;
+  std::vector<int32_t>& next = corpus.next;
+
+  ac2::CounterMap counts(1 << 16);
+  OccurrenceIndex index;
+  std::vector<std::pair<int32_t, int32_t>> merges;
+  std::vector<int32_t> token_cells(static_cast<size_t>(atom_count), 1);
+  int64_t initial_pair_types = 0;
+  int64_t heap_pops = 0;
+  int64_t counter_updates = 0;
+  int64_t compactions = 0;
+  int64_t peak_stored = 0;
+  int64_t oversized_pairs = 0;
+  int32_t observed_cells = 1;
+
+  {
+  py::gil_scoped_release release;
+
+  // One pass builds ascending occurrence lists; the counter is their length.
+  for (int64_t position = 0; position < initial_length; ++position) {
+    const int32_t follower = next[position];
+    if (follower < 0) continue;
+    index.append(ordered_key(symbol[position], symbol[follower]),
+                 static_cast<int32_t>(position));
+  }
+  for (size_t row = 0; row < index.lists.size(); ++row) {
+    counts.add(index.keys[row], static_cast<int64_t>(index.lists[row].size()));
+  }
+  initial_pair_types = static_cast<int64_t>(index.lists.size());
+  peak_stored = index.stored;
+
+  std::vector<ac2::HeapEntry> heap;
+  heap.reserve(index.lists.size());
+  counts.for_each([&](uint64_t key, int64_t value) { heap.push_back({value, key, 0}); });
+  std::make_heap(heap.begin(), heap.end(), ac2::HeapOrder{});
+
+  token_cells.reserve(static_cast<size_t>(vocab_size));
+  merges.reserve(static_cast<size_t>(vocab_size - atom_count));
+
+  ac2::CounterMap oversized(1 << 12);
+  StampSet touched(static_cast<size_t>(vocab_size));
+  std::vector<int32_t> occurrences;
+
+  while (static_cast<int64_t>(atom_count + merges.size()) < vocab_size && !heap.empty()) {
+    const ac2::HeapEntry top = ac2::heap_pop(heap);
+    ++heap_pops;
+    const uint64_t key = top.key;
+    const int64_t frequency = counts.get(key);
+    if (frequency <= 0 || frequency != top.count) continue;
+    if (oversized.get(key) != 0) continue;
+    const int32_t left = key_left(key);
+    const int32_t right = key_right(key);
+    const int64_t span = static_cast<int64_t>(token_cells[left]) + token_cells[right];
+    if (span > max_token_cells) {
+      oversized.add(key, 1);
+      continue;
+    }
+
+    const int32_t child = atom_count + static_cast<int32_t>(merges.size());
+    token_cells.push_back(static_cast<int32_t>(span));
+    observed_cells = std::max(observed_cells, static_cast<int32_t>(span));
+
+    occurrences.clear();
+    occurrences.swap(index.lists[static_cast<size_t>(index.slots.get(key) - 1)]);
+    index.stored -= static_cast<int64_t>(occurrences.size());
+    // Left-to-right application is what makes overlapping runs deterministic;
+    // later appends leave the list unsorted, so restore the order first.
+    std::sort(occurrences.begin(), occurrences.end());
+    touched.reset();
+    touched.insert(key);
+
+    int64_t applied = 0;
+    for (const int32_t position : occurrences) {
+      if (symbol[position] != left) continue;
+      const int32_t follower = next[position];
+      if (follower < 0 || symbol[follower] != right) continue;
+      const int32_t before = prev[position];
+      const int32_t after = next[follower];
+      if (before >= 0) {
+        const uint64_t stale = ordered_key(symbol[before], left);
+        counts.add(stale, -1);
+        touched.insert(stale);
+      }
+      if (after >= 0) {
+        const uint64_t stale = ordered_key(right, symbol[after]);
+        counts.add(stale, -1);
+        touched.insert(stale);
+      }
+      symbol[position] = child;
+      symbol[follower] = -1;
+      next[position] = after;
+      if (after >= 0) prev[after] = position;
+      prev[follower] = -1;
+      next[follower] = -1;
+      if (before >= 0) {
+        const uint64_t fresh = ordered_key(symbol[before], child);
+        counts.add(fresh, 1);
+        index.append(fresh, before);
+        touched.insert(fresh);
+      }
+      if (after >= 0) {
+        const uint64_t fresh = ordered_key(child, symbol[after]);
+        counts.add(fresh, 1);
+        index.append(fresh, position);
+        touched.insert(fresh);
+      }
+      ++applied;
+    }
+    counts.add(key, -applied);
+    counter_updates += 4 * applied;
+    length -= applied;
+    merges.emplace_back(left, right);
+
+    for (const uint64_t member : touched.members()) {
+      const int64_t value = counts.get(member);
+      if (value > 0) ac2::heap_push(heap, {value, member, 0});
+    }
+    peak_stored = std::max(peak_stored, index.stored);
+
+    // Superseded positions stay in the lists of pairs that never won.  Sweep
+    // them out once they outnumber the live corpus, so memory tracks the
+    // corpus rather than the merge history.
+    if (index.stored > 3 * length) {
+      int64_t retained = 0;
+      for (size_t row = 0; row < index.lists.size(); ++row) {
+        auto& list = index.lists[row];
+        if (list.empty()) continue;
+        const int32_t expected_left = key_left(index.keys[row]);
+        const int32_t expected_right = key_right(index.keys[row]);
+        size_t write = 0;
+        for (const int32_t position : list) {
+          const int32_t follower = next[position];
+          if (symbol[position] != expected_left) continue;
+          if (follower < 0 || symbol[follower] != expected_right) continue;
+          list[write++] = position;
+        }
+        list.resize(write);
+        if (write * 2 < list.capacity()) list.shrink_to_fit();
+        retained += static_cast<int64_t>(write);
+      }
+      index.stored = retained;
+      ++compactions;
+    }
+
+    if (heap.size() > 4 * counts.size() + 65536) {
+      heap.clear();
+      counts.for_each([&](uint64_t member, int64_t value) { heap.push_back({value, member, 0}); });
+      std::make_heap(heap.begin(), heap.end(), ac2::HeapOrder{});
+    }
+  }
+
+  oversized_pairs = static_cast<int64_t>(oversized.live_entries());
+  }
+
+  py::list merge_list;
+  for (const auto& merge : merges) {
+    merge_list.append(py::make_tuple(merge.first, merge.second));
+  }
+  py::dict report;
+  report["episodes"] = corpus.episodes;
+  report["actions"] = corpus.actions;
+  report["atoms"] = static_cast<int64_t>(atom_count);
+  report["merges"] = static_cast<int64_t>(merges.size());
+  report["vocab_size"] = static_cast<int64_t>(atom_count + merges.size());
+  report["cells"] = initial_length;
+  report["tokens_per_action_initial"] =
+      static_cast<double>(initial_length) / static_cast<double>(corpus.actions);
+  report["tokens_per_action"] =
+      static_cast<double>(length) / static_cast<double>(corpus.actions);
+  report["compression"] = static_cast<double>(initial_length) / static_cast<double>(length);
+  report["max_token_cells"] = max_token_cells;
+  report["max_token_cells_observed"] = static_cast<int64_t>(observed_cells);
+  report["oversized_pairs_skipped"] = oversized_pairs;
+  report["initial_pair_types"] = initial_pair_types;
+  report["pair_types_peak"] = static_cast<int64_t>(index.lists.size());
+  report["occurrence_slots_peak"] = peak_stored;
+  report["occurrence_compactions"] = compactions;
+  report["heap_pops"] = heap_pops;
+  report["counter_updates"] = counter_updates;
+  report["threads"] = thread_count;
+  report["requested_threads"] = threads;
+  report["trainer"] = "native_sparse_exact";
+  return py::make_tuple(merge_list, report);
+}

actioncodec2_kernels/common.hpp

@@ -0,0 +1,139 @@
+// Shared primitives for the ActionCodec2 native kernels.
+#pragma once
+
+#include <cstdint>
+#include <algorithm>
+#include <utility>
+#include <vector>
+
+namespace ac2 {
+
+class CounterMap {
+ public:
+  static constexpr uint64_t kEmpty = ~0ull;
+
+  explicit CounterMap(size_t capacity_hint = 1024) { reserve(capacity_hint); }
+
+  void reserve(size_t entries) {
+    size_t target = 16;
+    while (target * 3 < entries * 4) target <<= 1;  // keep load factor <= 0.75
+    if (target <= mask_ + 1 && !keys_.empty()) return;
+    rehash(target);
+  }
+
+  size_t size() const { return size_; }
+  size_t buckets() const { return mask_ + 1; }
+
+  void clear() {
+    std::fill(keys_.begin(), keys_.end(), kEmpty);
+    size_ = 0;
+  }
+
+  int64_t get(uint64_t key) const {
+    size_t slot = probe(key);
+    return keys_[slot] == key ? values_[slot] : 0;
+  }
+
+  int64_t add(uint64_t key, int64_t delta) {
+    return add_with_previous(key, delta).second;
+  }
+
+  // Update and expose both values with one hash-table probe. This is useful
+  // when a caller must distinguish a newly observed key from a stale one.
+  std::pair<int64_t, int64_t> add_with_previous(uint64_t key, int64_t delta) {
+    size_t slot = probe(key);
+    if (keys_[slot] != key) {
+      if ((size_ + 1) * 4 > (mask_ + 1) * 3) {
+        rehash((mask_ + 1) << 1);
+        slot = probe(key);
+      }
+      keys_[slot] = key;
+      values_[slot] = 0;
+      ++size_;
+    }
+    const int64_t previous = values_[slot];
+    values_[slot] += delta;
+    return {previous, values_[slot]};
+  }
+
+  template <typename Fn>
+  void for_each(Fn&& fn) const {
+    for (size_t slot = 0; slot <= mask_; ++slot) {
+      if (keys_[slot] != kEmpty && values_[slot] != 0) fn(keys_[slot], values_[slot]);
+    }
+  }
+
+  size_t live_entries() const {
+    size_t live = 0;
+    for (size_t slot = 0; slot <= mask_; ++slot) {
+      if (keys_[slot] != kEmpty && values_[slot] != 0) ++live;
+    }
+    return live;
+  }
+
+ private:
+  static uint64_t mix(uint64_t value) {  // splitmix64 finalizer
+    value ^= value >> 30;
+    value *= 0xbf58476d1ce4e5b9ull;
+    value ^= value >> 27;
+    value *= 0x94d049bb133111ebull;
+    return value ^ (value >> 31);
+  }
+
+  size_t probe(uint64_t key) const {
+    size_t slot = mix(key) & mask_;
+    while (keys_[slot] != kEmpty && keys_[slot] != key) slot = (slot + 1) & mask_;
+    return slot;
+  }
+
+  void rehash(size_t buckets) {
+    std::vector<uint64_t> old_keys;
+    std::vector<int64_t> old_values;
+    old_keys.swap(keys_);
+    old_values.swap(values_);
+    keys_.assign(buckets, kEmpty);
+    values_.assign(buckets, 0);
+    mask_ = buckets - 1;
+    size_ = 0;
+    for (size_t slot = 0; slot < old_keys.size(); ++slot) {
+      if (old_keys[slot] != kEmpty && old_values[slot] != 0) {
+        size_t target = probe(old_keys[slot]);
+        keys_[target] = old_keys[slot];
+        values_[target] = old_values[slot];
+        ++size_;
+      }
+    }
+  }
+
+  std::vector<uint64_t> keys_;
+  std::vector<int64_t> values_;
+  size_t mask_ = 0;
+  size_t size_ = 0;
+};
+
+struct HeapEntry {
+  int64_t count;
+  uint64_t key;
+  uint32_t version;  // only shifted temporal self-rules use this
+};
+
+struct HeapOrder {
+  bool operator()(const HeapEntry& a, const HeapEntry& b) const {
+    if (a.count != b.count) return a.count < b.count;  // larger count first
+    return a.key > b.key;                              // smaller key first
+  }
+};
+
+inline void heap_push(std::vector<HeapEntry>& heap, const HeapEntry& entry) {
+  heap.push_back(entry);
+  std::push_heap(heap.begin(), heap.end(), HeapOrder{});
+}
+
+inline HeapEntry heap_pop(std::vector<HeapEntry>& heap) {
+  std::pop_heap(heap.begin(), heap.end(), HeapOrder{});
+  HeapEntry top = heap.back();
+  heap.pop_back();
+  return top;
+}
+
+}  // namespace ac2

actioncodec2_kernels/grammar.cpp

@@ -1,162 +1,279 @@
+// Coverage grammars and multi-component routing, including time-major streams.
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
-
 #include <algorithm>
 #include <cstdint>
-#include <limits>
-#include <map>
 #include <memory>
 #include <stdexcept>
 #include <vector>
+#include "parallel.hpp"
+#include "token_streams.hpp"
 
-namespace py = pybind11;
+namespace py=pybind11;
 namespace {
-using State = std::vector<int64_t>;
-
-class CoverageGrammar {
- protected:
-  int64_t horizon_, vocabulary_;
-  virtual void validate_state(const State&) const = 0;
-  virtual bool allowed(const State&, int64_t) const = 0;
-  virtual State apply(const State&, int64_t) const = 0;
+using State=std::vector<int64_t>;
+using States=std::vector<State>;
+class Grammar {
  public:
-  virtual ~CoverageGrammar() = default;
-  State advance(const State& state, int64_t token) const {
-    validate_state(state);
-    if (token < 0 || token >= vocabulary_ || !allowed(state, token)) return {};
-    return apply(state, token);
-  }
-  py::array_t<bool> mask(const State& state) const {
-    validate_state(state);
-    py::array_t<bool> result(vocabulary_);
-    auto* output = result.mutable_data();
+  int64_t horizon,vocabulary;
+  virtual ~Grammar()=default;
+  virtual State initial()const=0;
+  virtual State advance(const State&,int64_t)const=0;
+  virtual int64_t covered(const State&)const=0;
+  virtual int64_t next_dimension(const State&)const=0;
+  virtual bool allowed(const State&,int64_t)const=0;
+  virtual std::shared_ptr<Grammar> clone(int64_t)const=0;
+  py::array_t<bool> mask(const State& state)const {
+    py::array_t<bool> result(vocabulary);
     py::gil_scoped_release release;
-    for (int64_t token = 0; token < vocabulary_; ++token) output[token] = allowed(state, token);
+    for(int64_t token=0;token<vocabulary;++token)result.mutable_data()[token]=allowed(state,token);
     return result;
   }
 };
-
-class SetGrammar : public CoverageGrammar {
+class SetGrammar:public Grammar {
   uint64_t full_;
-  std::vector<int64_t> roots_;
-  std::vector<std::vector<std::pair<int64_t, uint64_t>>> patterns_;
-  std::vector<bool> scope_;
-
-  void validate_state(const State& state) const override {
-    if (static_cast<int64_t>(state.size()) != horizon_)
-      throw std::invalid_argument("invalid grammar state length");
-    for (int64_t value : state) {
-      if (value < 0 || (static_cast<uint64_t>(value) & ~full_))
-        throw std::invalid_argument("invalid grammar state support");
+  using Patterns=std::vector<std::vector<std::pair<int64_t,uint64_t>>>;
+  std::shared_ptr<const std::vector<int64_t>> roots_;
+  std::shared_ptr<const Patterns> patterns_;
+  std::shared_ptr<const std::vector<bool>> scope_;
+ public:
+  SetGrammar(int64_t steps,uint64_t full,std::vector<int64_t> roots,
+             std::vector<std::vector<std::pair<int64_t,uint64_t>>> patterns)
+      :full_(full),roots_(std::make_shared<const std::vector<int64_t>>(std::move(roots))),
+       patterns_(std::make_shared<const Patterns>(std::move(patterns))) {
+    if(steps<1 || !full || roots_->size()!=patterns_->size())throw std::invalid_argument("invalid grammar metadata");
+    horizon=steps;vocabulary=roots_->size();
+    std::vector<bool> scopes;
+    for(const auto& pattern:*patterns_) {
+      bool scope=true;
+      for(auto[offset,support]:pattern) {
+        if(offset<0)throw std::invalid_argument("negative grammar pattern offset");
+        scope=scope && !(support & ~full_);
+      }
+      scopes.push_back(scope);
     }
+    scope_=std::make_shared<const std::vector<bool>>(std::move(scopes));
+  }
+  std::shared_ptr<Grammar> clone(int64_t steps)const override {
+    if(steps<1)throw std::invalid_argument("horizon must be positive");
+    auto result=std::make_shared<SetGrammar>(*this);result->horizon=steps;return result;
+  }
+  State initial()const override{return State(horizon,0);}
+  int64_t covered(const State& state)const override {
+    if(static_cast<int64_t>(state.size())!=horizon)throw std::invalid_argument("invalid grammar state");
+    for(int64_t i=0;i<horizon;++i)if(static_cast<uint64_t>(state[i])!=full_)return i;
+    return horizon;
   }
-  int64_t front(const State& state) const {
-    for (int64_t t = 0; t < horizon_; ++t)
-      if (static_cast<uint64_t>(state[t]) != full_) return t;
-    return horizon_;
-  }
-  bool allowed(const State& state, int64_t token) const override {
-    if (!scope_[token]) return false;
-    const int64_t start = front(state);
-    if (start == horizon_) return false;
-    uint64_t missing = full_ & ~static_cast<uint64_t>(state[start]);
-    int64_t root = 0;
-    while (!(missing & 1)) { missing >>= 1; ++root; }
-    if (roots_[token] != root) return false;
-    for (auto [offset, support] : patterns_[token]) {
-      if (offset >= horizon_ - start || (static_cast<uint64_t>(state[start + offset]) & support))
-        return false;
+  int64_t next_dimension(const State& state)const override {
+    int64_t front=covered(state);if(front==horizon)return -1;
+    uint64_t missing=full_ & ~static_cast<uint64_t>(state[front]);
+    int64_t d=0;while(!(missing & 1)){missing>>=1;++d;}return d;
+  }
+  bool allowed(const State& state,int64_t token)const override {
+    if(token<0 || token>=vocabulary || !(*scope_)[token])return false;
+    const int64_t front=covered(state);
+    if(front==horizon || (*roots_)[token]!=next_dimension(state))return false;
+    for(auto[offset,support]:(*patterns_)[token]) {
+      if(front+offset>=horizon || (static_cast<uint64_t>(state[front+offset]) & support))return false;
     }
     return true;
   }
-  State apply(const State& state, int64_t token) const override {
-    State result = state;
-    const int64_t start = front(state);
-    for (auto [offset, support] : patterns_[token])
-      result[start + offset] |= static_cast<int64_t>(support);
+  State advance(const State& state,int64_t token)const override {
+    if(!allowed(state,token))return {};
+    State result=state;int64_t front=covered(state);
+    for(auto[offset,support]:(*patterns_)[token])result[front+offset]|=static_cast<int64_t>(support);
     return result;
   }
+};
+class PlainGrammar:public Grammar {
+  std::shared_ptr<const std::vector<int64_t>> lengths_,dims_;
+  std::shared_ptr<const std::vector<bool>> phase_;
  public:
-  SetGrammar(int64_t horizon, uint64_t full, std::vector<int64_t> roots,
-             std::vector<std::vector<std::pair<int64_t, uint64_t>>> patterns)
-      : full_(full), roots_(std::move(roots)), patterns_(std::move(patterns)) {
-    if (horizon < 1 || !full || full > static_cast<uint64_t>(INT64_MAX) ||
-        roots_.empty() || roots_.size() != patterns_.size())
-      throw std::invalid_argument("invalid Set-BPE grammar metadata");
-    horizon_ = horizon;
-    vocabulary_ = roots_.size();
-    for (int64_t token = 0; token < vocabulary_; ++token) {
-      if (roots_[token] < 0 || roots_[token] >= 63 || patterns_[token].empty())
-        throw std::invalid_argument("invalid grammar root or empty pattern");
-      std::map<int64_t, uint64_t> coverage;
-      bool scope = true;
-      for (auto [offset, support] : patterns_[token]) {
-        if (offset < 0 || !support || support > static_cast<uint64_t>(INT64_MAX) ||
-            (coverage[offset] & support))
-          throw std::invalid_argument("invalid grammar pattern offset or support");
-        coverage[offset] |= support;
-        scope = scope && !(support & ~full_);
-      }
-      const uint64_t anchor = coverage[0];
-      if (!(anchor & (uint64_t{1} << roots_[token])) ||
-          (anchor & ((uint64_t{1} << roots_[token]) - 1)))
-        throw std::invalid_argument("grammar root must anchor the pattern");
-      scope_.push_back(scope);
-    }
+  PlainGrammar(int64_t steps,std::vector<int64_t> lengths,std::vector<int64_t> dims,
+                py::array_t<bool,py::array::c_style|py::array::forcecast> phase)
+      :lengths_(std::make_shared<const std::vector<int64_t>>(std::move(lengths))),
+       dims_(std::make_shared<const std::vector<int64_t>>(std::move(dims))) {
+    horizon=steps;vocabulary=lengths_->size();
+    if(steps<1 || dims_->empty() || phase.ndim()!=2 || phase.shape(0)!=static_cast<int64_t>(dims_->size()) || phase.shape(1)!=vocabulary)
+      throw std::invalid_argument("invalid plain-BPE grammar metadata");
+    for(auto value:*lengths_)if(value<1)throw std::invalid_argument("invalid token expansion");
+    phase_=std::make_shared<const std::vector<bool>>(phase.data(),phase.data()+phase.size());
+  }
+  std::shared_ptr<Grammar> clone(int64_t steps)const override {
+    if(steps<1)throw std::invalid_argument("horizon must be positive");
+    auto result=std::make_shared<PlainGrammar>(*this);result->horizon=steps;return result;
+  }
+  State initial()const override{return {0};}
+  int64_t covered(const State& state)const override {
+    if(state.size()!=1 || state[0]<0)throw std::invalid_argument("invalid grammar state");
+    return state[0]/dims_->size();
+  }
+  int64_t next_dimension(const State& state)const override{return (*dims_)[state[0]%dims_->size()];}
+  bool allowed(const State& state,int64_t token)const override {
+    return token>=0 && token<vocabulary && state[0]+(*lengths_)[token]<=horizon*static_cast<int64_t>(dims_->size()) && (*phase_)[(state[0]%dims_->size())*vocabulary+token];
+  }
+  State advance(const State& state,int64_t token)const override {
+    if(!allowed(state,token))return {};
+    return {state[0]+(*lengths_)[token]};
   }
 };
-
-class PlainGrammar : public CoverageGrammar {
-  int64_t width_, cells_;
-  std::vector<int64_t> lengths_;
-  std::vector<bool> phase_;
-
-  void validate_state(const State& state) const override {
-    if (state.size() != 1 || state[0] < 0 || state[0] > cells_)
-      throw std::invalid_argument("invalid plain-BPE grammar state");
+using Grammars=std::vector<std::shared_ptr<Grammar>>;
+int64_t next_component(const Grammars& grammars,const States& states,bool time_major) {
+  int64_t selected=grammars.size(),minimum=INT64_MAX;
+  for(size_t i=0;i<grammars.size();++i) {
+    const int64_t front=grammars[i]->covered(states[i]);
+    if(front>=grammars[i]->horizon)continue;
+    if(!time_major)return i;
+    if(front<minimum){minimum=front;selected=i;}
   }
-  bool allowed(const State& state, int64_t token) const override {
-    return lengths_[token] <= cells_ - state[0] &&
-           phase_[(state[0] % width_) * vocabulary_ + token];
+  return selected;
+}
+States initial(const Grammars& grammars) {
+  States states;for(const auto& grammar:grammars)states.push_back(grammar->initial());return states;
+}
+bool advance(const Grammars& grammars,const std::vector<int64_t>& offsets,States& states,int64_t index,int64_t token) {
+  if(index<0 || index>=static_cast<int64_t>(grammars.size()))return false;
+  State next=grammars[index]->advance(states[index],token-offsets[index]);
+  if(next.empty())return false;
+  states[index]=std::move(next);
+  return true;
+}
+py::tuple state_tuple(const State& state) {
+  py::tuple result(state.size());for(size_t i=0;i<state.size();++i)result[i]=state[i];return result;
+}
+py::object route_advance(const Grammars& grammars,const std::vector<int64_t>& offsets,
+                         bool time_major,int64_t index,States states,int64_t token) {
+  if(states.size()!=grammars.size() || offsets.size()!=grammars.size())throw std::invalid_argument("invalid routed metadata");
+  if(!advance(grammars,offsets,states,index,token))return py::none();
+  py::tuple local(states.size());for(size_t i=0;i<states.size();++i)local[i]=state_tuple(states[i]);
+  return py::make_tuple(next_component(grammars,states,time_major),local);
+}
+std::vector<int64_t> interleave(const Grammars& grammars,const std::vector<int64_t>& offsets,
+                                bool time_major,const std::vector<std::vector<int64_t>>& streams) {
+  if(streams.size()!=grammars.size() || offsets.size()!=grammars.size())throw std::invalid_argument("one stream per component is required");
+  States states=initial(grammars);std::vector<size_t> cursor(streams.size());std::vector<int64_t> output;
+  int64_t index;
+  while((index=next_component(grammars,states,time_major))<static_cast<int64_t>(grammars.size())) {
+    if(cursor[index]==streams[index].size())throw std::invalid_argument("component stream ends before complete coverage");
+    const int64_t token=streams[index][cursor[index]++]+offsets[index];
+    if(!advance(grammars,offsets,states,index,token))throw std::invalid_argument("invalid routed component stream");
+    output.push_back(token);
   }
-  State apply(const State& state, int64_t token) const override {
-    return {state[0] + lengths_[token]};
+  for(size_t i=0;i<cursor.size();++i)if(cursor[i]!=streams[i].size())throw std::invalid_argument("extra component tokens");
+  return output;
+}
+std::vector<std::vector<int64_t>> split(const Grammars& grammars,const std::vector<int64_t>& offsets,
+                                       bool time_major,const std::vector<int64_t>& tokens) {
+  if(offsets.size()!=grammars.size())throw std::invalid_argument("invalid routed offsets");
+  States states=initial(grammars);std::vector<std::vector<int64_t>> output(grammars.size());
+  for(int64_t token:tokens) {
+    int64_t index=next_component(grammars,states,time_major);
+    if(!advance(grammars,offsets,states,index,token))throw std::invalid_argument("invalid routed token coverage");
+    output[index].push_back(token-offsets[index]);
   }
- public:
-  PlainGrammar(int64_t horizon, std::vector<int64_t> lengths,
-               const py::array_t<bool, py::array::c_style | py::array::forcecast>& phase)
-      : lengths_(std::move(lengths)) {
-    if (horizon < 1 || lengths_.empty() || phase.ndim() != 2 || phase.shape(0) < 1 ||
-        phase.shape(1) != static_cast<int64_t>(lengths_.size()) ||
-        horizon > std::numeric_limits<int64_t>::max() / phase.shape(0))
-      throw std::invalid_argument("invalid plain-BPE grammar metadata");
-    for (int64_t length : lengths_) {
-      if (length < 1) throw std::invalid_argument("token expansion must be positive");
+  if(next_component(grammars,states,time_major)!=static_cast<int64_t>(grammars.size()))
+    throw std::invalid_argument("Incomplete action token stream");
+  return output;
+}
+
+using Streams=std::vector<std::vector<int64_t>>;
+using ComponentStreams=std::vector<Streams>;
+
+Streams interleave_batch(const Grammars& grammars,const std::vector<int64_t>& offsets,
+                        bool time_major,const ComponentStreams& streams) {
+  if(streams.size()!=grammars.size() || streams.empty() || streams.front().empty())
+    throw std::invalid_argument("one non-empty batch per component is required");
+  const int64_t batch=streams.front().size();
+  for(const auto& rows:streams)if(static_cast<int64_t>(rows.size())!=batch)
+    throw std::invalid_argument("component batches must have the same size");
+  Streams output(batch);
+  py::gil_scoped_release release;
+  ac2::parallel_for(batch,0,[&](int64_t b) {
+    Streams local;for(const auto& rows:streams)local.push_back(rows[b]);
+    output[b]=interleave(grammars,offsets,time_major,local);
+  });
+  return output;
+}
+
+py::tuple decode_batch(const Grammars& prototypes,const std::vector<int64_t>& offsets,
+                       bool time_major,const py::iterable& streams,
+                       py::array_t<int64_t,py::array::c_style|py::array::forcecast> counts,
+                       int64_t cells_per_step) {
+  if(prototypes.empty() || counts.ndim()!=1 || cells_per_step<1)
+    throw std::invalid_argument("invalid routed batch metadata");
+  const auto rows=ac2::checked_token_streams(streams,counts.size(),true);
+  int64_t steps=-1;
+  ComponentStreams output(prototypes.size(),Streams(rows.size()));
+  {
+    py::gil_scoped_release release;
+    for(const auto& row:rows) {
+      int64_t cells=0;
+      for(int64_t token:row) {
+        if(counts.data()[token]<1)throw std::invalid_argument("token is outside the routed profile vocabularies");
+        cells+=counts.data()[token];
+      }
+      if(cells%cells_per_step || cells/cells_per_step<1)
+        throw std::invalid_argument("Incomplete action token stream: " +
+            std::to_string(cells_per_step-cells%cells_per_step) +
+            " more coordinate values are needed to finish a step. "
+            "Constrain generation with codec.grammar(H, fps=...).next_token_mask(prefix) "
+            "and check is_complete(tokens) for the intended H before decoding. "
+            "Use codec.batch_decode(..., errors='return') to isolate invalid rows; "
+            "do not pad truncated tokens with zeros.");
+      if(steps>=0 && steps!=cells/cells_per_step)
+        throw std::invalid_argument("all token streams must decode to the same length; use codec.batch_decode(...) for independent variable-length outputs");
+      steps=cells/cells_per_step;
     }
-    horizon_ = horizon;
-    width_ = phase.shape(0);
-    cells_ = horizon * width_;
-    vocabulary_ = lengths_.size();
-    phase_.assign(phase.data(), phase.data() + phase.size());
+    Grammars grammars;for(const auto& prototype:prototypes)grammars.push_back(prototype->clone(steps));
+    ac2::parallel_for(rows.size(),0,[&](int64_t b) {
+      auto local=split(grammars,offsets,time_major,rows[b]);
+      for(size_t i=0;i<local.size();++i)output[i][b]=std::move(local[i]);
+    });
   }
-};
+  return py::make_tuple(steps,output);
+}
+
+std::vector<int64_t> complete(const Grammars& grammars,const std::vector<int64_t>& offsets,
+                              bool time_major,std::vector<int64_t> tokens,int64_t steps,
+                              const std::vector<std::vector<int64_t>>& zero_atoms) {
+  if(grammars.empty() || offsets.size()!=grammars.size() || zero_atoms.size()!=grammars.size() ||
+     steps<1 || steps>grammars.front()->horizon)throw std::invalid_argument("invalid prefix completion metadata");
+  py::gil_scoped_release release;
+  States states=initial(grammars);
+  for(int64_t token:tokens) {
+    int64_t index=next_component(grammars,states,time_major);
+    if(!advance(grammars,offsets,states,index,token))throw std::invalid_argument("invalid routed token prefix");
+  }
+  for(size_t i=0;i<grammars.size();++i)if(grammars[i]->covered(states[i])<steps)
+    throw std::invalid_argument("action prefix does not cover the requested frames");
+  int64_t index;
+  while((index=next_component(grammars,states,time_major))<static_cast<int64_t>(grammars.size())) {
+    int64_t dimension=grammars[index]->next_dimension(states[index]);
+    if(dimension<0 || dimension>=static_cast<int64_t>(zero_atoms[index].size()))
+      throw std::invalid_argument("missing zero-motion primitive atom");
+    const int64_t token=offsets[index]+zero_atoms[index][dimension];
+    if(!advance(grammars,offsets,states,index,token))throw std::invalid_argument("invalid zero-motion suffix");
+    tokens.push_back(token);
+  }
+  return tokens;
 }
+}  // namespace
 
 void register_grammar(py::module_& module) {
-  py::class_<CoverageGrammar, std::shared_ptr<CoverageGrammar>>(module, "CoverageGrammar")
-      .def("advance", [](const CoverageGrammar& grammar, const State& state, int64_t token) -> py::object {
-        State next = grammar.advance(state, token);
-        if (next.empty()) return py::none();
-        py::tuple result(next.size());
-        for (size_t i = 0; i < next.size(); ++i) result[i] = next[i];
-        return result;
+  py::class_<Grammar,std::shared_ptr<Grammar>>(module,"CoverageGrammar")
+      .def("advance",[](const Grammar& g,const State& s,int64_t token)->py::object {
+        State next=g.advance(s,token);return next.empty()?py::object(py::none()):py::object(state_tuple(next));
       })
-      .def("mask", &CoverageGrammar::mask);
-  py::class_<SetGrammar, CoverageGrammar, std::shared_ptr<SetGrammar>>(module, "SetGrammar")
-      .def(py::init<int64_t, uint64_t, std::vector<int64_t>,
-                    std::vector<std::vector<std::pair<int64_t, uint64_t>>>>());
-  py::class_<PlainGrammar, CoverageGrammar, std::shared_ptr<PlainGrammar>>(module, "PlainGrammar")
-      .def(py::init<int64_t, std::vector<int64_t>,
-                    const py::array_t<bool, py::array::c_style | py::array::forcecast>&>());
+      .def("mask",&Grammar::mask).def("covered",&Grammar::covered).def("clone",&Grammar::clone);
+  py::class_<SetGrammar,Grammar,std::shared_ptr<SetGrammar>>(module,"SetGrammar")
+      .def(py::init<int64_t,uint64_t,std::vector<int64_t>,std::vector<std::vector<std::pair<int64_t,uint64_t>>>>());
+  py::class_<PlainGrammar,Grammar,std::shared_ptr<PlainGrammar>>(module,"PlainGrammar")
+      .def(py::init<int64_t,std::vector<int64_t>,std::vector<int64_t>,py::array_t<bool,py::array::c_style|py::array::forcecast>>());
+  module.def("route_advance",&route_advance);
+  module.def("route_interleave",&interleave,py::call_guard<py::gil_scoped_release>());
+  module.def("route_split",&split,py::call_guard<py::gil_scoped_release>());
+  module.def("route_interleave_batch",&interleave_batch);
+  module.def("route_decode_batch",&decode_batch);
+  module.def("route_complete",&complete);
 }

actioncodec2_kernels/inference.cpp

@@ -0,0 +1,245 @@
+// Lossless plain-BPE inference and additive physical reconstruction.
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <algorithm>
+#include <cstdint>
+#include <map>
+#include <limits>
+#include <queue>
+#include <stdexcept>
+#include <vector>
+#include "parallel.hpp"
+#include "quaternion.hpp"
+
+namespace py = pybind11;
+namespace {
+using I64 = py::array_t<int64_t, py::array::c_style | py::array::forcecast>;
+using F64 = py::array_t<double, py::array::c_style | py::array::forcecast>;
+using F32 = py::array_t<float, py::array::c_style | py::array::forcecast>;
+
+
+
+F64 quaternion_map(const F64& values, const std::string& operation) {
+  const bool vector_in = operation == "from_rotvec" || operation == "from_euler";
+  const int input_width=vector_in?3:4;
+  const bool vector_out=operation == "to_rotvec" || operation == "to_euler";
+  const int width=vector_out?3:4;
+  if(values.ndim()!=2 || values.shape(1)!=input_width)
+    throw std::invalid_argument("invalid quaternion operation shape");
+  F64 output({values.shape(0),static_cast<int64_t>(width)});
+  py::gil_scoped_release release;
+  for(int64_t i=0;i<values.shape(0);++i) {
+    const double* v=values.data()+i*input_width;
+    ac2q::Q q=vector_in?(operation=="from_rotvec"?ac2q::from_rotvec(v):ac2q::from_euler(v))
+                         :ac2q::Q{v[0],v[1],v[2],v[3]};
+    double* out=output.mutable_data()+i*width;
+    if(operation=="to_rotvec")ac2q::to_rotvec(q,out);
+    else if(operation=="to_euler")ac2q::to_euler(q,out);
+    else {
+      if(operation=="normalize")q=ac2q::normalize(q);
+      else if(operation=="conjugate")q=ac2q::conjugate(q);
+      else if(!vector_in)throw std::invalid_argument("unknown quaternion operation");
+      std::copy(q.begin(),q.end(),out);
+    }
+  }
+  return output;
+}
+
+F64 quaternion_multiply(const F64& first, const F64& second) {
+  if(first.ndim()!=2 || first.shape(1)!=4 || second.ndim()!=2 ||
+     second.shape(1)!=4 || first.shape(0)!=second.shape(0))
+    throw std::invalid_argument("invalid quaternion product shape");
+  F64 output({first.shape(0),int64_t{4}});
+  py::gil_scoped_release release;
+  for(int64_t i=0;i<first.shape(0);++i) {
+    auto* a=first.data()+4*i;auto* b=second.data()+4*i;
+    auto q=ac2q::multiply({a[0],a[1],a[2],a[3]},{b[0],b[1],b[2],b[3]});
+    std::copy(q.begin(),q.end(),output.mutable_data()+4*i);
+  }
+  return output;
+}
+
+F32 so3_reconstruct(const F64& table, const I64& cells, const F32& previous,
+                    const F32& scale, int input_order, int primitive_order,
+                    bool left, bool rotation_vector, double radius) {
+  if(cells.ndim()!=3 || cells.shape(2)!=3 || table.ndim()!=2 || table.shape(0)!=3 ||
+     previous.ndim()!=2 || previous.shape(0)!=cells.shape(0) || previous.shape(1)!=3 || scale.size()!=3)
+    throw std::invalid_argument("invalid SO(3) reconstruction shapes");
+  const int64_t batch=cells.shape(0),steps=cells.shape(1);
+  F32 output({batch,steps,int64_t{3}});
+  py::gil_scoped_release release;
+  ac2::parallel_for(batch,0,[&](int64_t b) {
+    double boundary[3];
+    for(int d=0;d<3;++d)boundary[d]=static_cast<double>(previous.data()[3*b+d])*scale.data()[d];
+    auto increment=rotation_vector?ac2q::from_rotvec(boundary):ac2q::from_euler(boundary);
+    ac2q::Q state{0,0,0,1};
+    float accumulated[3]={0,0,0};
+    auto apply=[&](ac2q::Q a,ac2q::Q v){return left?ac2q::multiply(v,a):ac2q::multiply(a,v);};
+    for(int64_t t=0;t<steps;++t) {
+      double v[3],represented[3];
+      for(int d=0;d<3;++d) {
+        const int64_t cell=cells.data()[(b*steps+t)*3+d];
+        if(cell<0 || cell>=table.shape(1))throw std::invalid_argument("invalid primitive cell");
+        v[d]=table.data()[d*table.shape(1)+cell];
+      }
+      if(primitive_order==2) {
+        const double norm=std::sqrt(v[0]*v[0]+v[1]*v[1]+v[2]*v[2]);
+        const double factor=std::min(1.,radius/std::max(norm,std::numeric_limits<double>::min()));
+        for(double& x:v)x*=factor;
+        increment=apply(increment,ac2q::from_rotvec(v));
+        state=apply(state,increment);
+        auto q=input_order==1?increment:state;
+        if(rotation_vector)ac2q::to_rotvec(q,represented);else ac2q::to_euler(q,represented);
+      } else if(input_order==1) {
+        if(rotation_vector)std::copy(v,v+3,represented);
+        else {for(double& x:v)x=static_cast<float>(x);ac2q::to_euler(ac2q::from_rotvec(v),represented);}
+      } else {
+        state=apply(state,ac2q::from_rotvec(v));
+        if(rotation_vector)ac2q::to_rotvec(state,represented);else ac2q::to_euler(state,represented);
+      }
+      for(int d=0;d<3;++d) {
+        float value=static_cast<float>(represented[d]/scale.data()[d]);
+        if(primitive_order==1 && input_order==1) {
+          const float next=accumulated[d]+value;value=next-accumulated[d];accumulated[d]=next;
+        }
+        output.mutable_data()[(b*steps+t)*3+d]=value;
+      }
+    }
+  });
+  return output;
+}
+
+F32 additive_reconstruct(const F64& table, const I64& cells, const F32& previous,
+                         const F32& scale, int input_order, int primitive_order) {
+  if (cells.ndim() != 3 || table.ndim() != 2 || scale.ndim() != 1 ||
+      cells.shape(2) != table.shape(0) || scale.size() != cells.shape(2) ||
+      previous.ndim() != 2 || previous.shape(0) != cells.shape(0) || previous.shape(1) != cells.shape(2))
+    throw std::invalid_argument("invalid additive reconstruction shapes");
+  if (primitive_order != 2 || input_order < 0 || input_order > 2)
+    throw std::invalid_argument("invalid differential order");
+  const int64_t batch = cells.shape(0), steps = cells.shape(1), width = cells.shape(2);
+  F32 output({batch, steps, width});
+  py::gil_scoped_release release;
+  ac2::parallel_for(batch, 0, [&](int64_t b) {
+    for (int64_t d = 0; d < width; ++d) {
+      double change_sum = 0, position = 0;
+      const float boundary = previous.data()[b * width + d] * scale.data()[d];
+      for (int64_t t = 0; t < steps; ++t) {
+        const int64_t index = (b * steps + t) * width + d, cell = cells.data()[index];
+        if (cell < 0 || cell >= table.shape(1)) throw std::invalid_argument("invalid cell");
+        const double primitive = table.data()[d * table.shape(1) + cell];
+        change_sum += primitive;
+        const double velocity = change_sum + boundary;
+        position += velocity;
+        const double value = input_order == 2 ? primitive : input_order == 1 ? velocity : position;
+        output.mutable_data()[index] = static_cast<float>(value / scale.data()[d]);
+      }
+    }
+  });
+  return output;
+}
+
+I64 choose_cells(const F32& table,const I64& bins,const F32& request,
+                  const F32& direction,const F32& limits,bool smooth) {
+  if(table.ndim()!=2 || request.ndim()!=2 || direction.ndim()!=2 ||
+     direction.shape(0)!=request.shape(0) || direction.shape(1)!=request.shape(1) ||
+     table.shape(0)!=request.shape(1) || bins.size()!=request.shape(1) || limits.size()!=request.shape(1))
+    throw std::invalid_argument("invalid primitive selection shapes");
+  I64 output({request.shape(0),request.shape(1)});
+  py::gil_scoped_release release;
+  const int64_t width=request.shape(1);
+  for(int64_t i=0;i<request.shape(0);++i)for(int64_t d=0;d<width;++d) {
+    int64_t nearest=0,feasible=-1,preferred=-1;
+    float best=INFINITY,best_feasible=INFINITY,best_preferred=INFINITY;
+    const float q=request.data()[i*width+d],sign=direction.data()[i*width+d];
+    for(int64_t j=0;j<bins.data()[d];++j) {
+      float p=table.data()[d*table.shape(1)+j],distance=std::abs(q-p);
+      if(distance<best){best=distance;nearest=j;}
+      if(distance<=limits.data()[d]+8*std::numeric_limits<float>::epsilon()) {
+        if(distance<best_feasible){best_feasible=distance;feasible=j;}
+        if((p>0?1.F:p<0?-1.F:0.F)*sign>=0 && distance<best_preferred) {
+          best_preferred=distance;preferred=j;
+        }
+      }
+    }
+    output.mutable_data()[i*width+d]=smooth?(preferred>=0?preferred:feasible>=0?feasible:nearest):nearest;
+  }
+  return output;
+}
+
+F32 additive_decode(const F32& table,const I64& cells,const F32& scale,
+                     const F32& lower,const F32& upper,int input_order,int primitive_order) {
+  if(cells.ndim()!=3 || table.ndim()!=2 || table.shape(0)!=cells.shape(2) ||
+     scale.size()!=cells.shape(2) || lower.size()!=cells.shape(2) || upper.size()!=cells.shape(2) ||
+     primitive_order<0 || primitive_order>1 || input_order<0 || input_order>2)
+    throw std::invalid_argument("invalid additive decoding metadata");
+  const int64_t batch=cells.shape(0),steps=cells.shape(1),width=cells.shape(2);
+  F32 output({batch,steps,width});
+  py::gil_scoped_release release;
+  ac2::parallel_for(batch,0,[&](int64_t b) {
+    for(int64_t d=0;d<width;++d) {
+      float position=0;
+      for(int64_t t=0;t<steps;++t) {
+        const int64_t index=(b*steps+t)*width+d,cell=cells.data()[index];
+        if(cell<0 || cell>=table.shape(1))throw std::invalid_argument("invalid primitive cell");
+        const float picked=table.data()[d*table.shape(1)+cell];
+        position=primitive_order?std::clamp(position+picked,lower.data()[d],upper.data()[d]):picked;
+        output.mutable_data()[index]=position;
+      }
+      for(int order=0;order<input_order;++order)for(int64_t t=steps-1;t>0;--t)
+        output.mutable_data()[(b*steps+t)*width+d]-=output.data()[(b*steps+t-1)*width+d];
+      for(int64_t t=0;t<steps;++t)output.mutable_data()[(b*steps+t)*width+d]/=scale.data()[d];
+    }
+  });
+  return output;
+}
+
+I64 spatial_expand(const std::vector<int64_t>& tokens,const I64& expansions,
+                    const I64& supports,uint64_t full) {
+  if(expansions.ndim()!=2 || supports.ndim()!=1 || expansions.shape(0)!=supports.size())
+    throw std::invalid_argument("invalid spatial metadata");
+  const int64_t width=expansions.shape(1);I64 output(width);
+  std::fill(output.mutable_data(),output.mutable_data()+width,-1);
+  uint64_t covered=0;
+  py::gil_scoped_release release;
+  for(int64_t token:tokens) {
+    if(token<0 || token>=supports.size())throw std::invalid_argument("token is outside the vocabulary");
+    const uint64_t support=supports.data()[token];
+    if(covered & support)throw std::invalid_argument("token supports overlap");
+    covered|=support;
+    for(int64_t d=0;d<width;++d) {
+      const int64_t value=expansions.data()[token*width+d];
+      if(value>=0)output.mutable_data()[d]=value;
+    }
+  }
+  if(covered!=full)throw std::invalid_argument("token set does not exactly cover the requested active dimensions");
+  return output;
+}
+F32 additive_transform(const F32& values,int order,bool integrate) {
+  if(values.ndim()!=3 || order<0 || order>2)throw std::invalid_argument("expected [B,T,D] and order 0..2");
+  const int64_t batch=values.shape(0),steps=values.shape(1),width=values.shape(2);
+  F32 output({batch,steps,width});std::copy(values.data(),values.data()+values.size(),output.mutable_data());
+  py::gil_scoped_release release;
+  for(int level=0;level<order;++level)for(int64_t b=0;b<batch;++b)for(int64_t d=0;d<width;++d) {
+    if(integrate) {
+      float sum=0;
+      for(int64_t t=0;t<steps;++t) {const int64_t i=(b*steps+t)*width+d;sum+=output.data()[i];output.mutable_data()[i]=sum;}
+    } else {
+      for(int64_t t=steps-1;t>0;--t)output.mutable_data()[(b*steps+t)*width+d]-=output.data()[(b*steps+t-1)*width+d];
+    }
+  }
+  return output;
+}
+}  // namespace
+
+void register_inference(py::module_& module) {
+  module.def("additive_reconstruct", &additive_reconstruct);
+  module.def("so3_reconstruct", &so3_reconstruct);
+  module.def("quaternion_map", &quaternion_map);
+  module.def("quaternion_multiply", &quaternion_multiply);
+  module.def("choose_cells", &choose_cells);
+  module.def("additive_decode", &additive_decode);
+  module.def("spatial_expand", &spatial_expand);
+  module.def("additive_transform", &additive_transform);
+}

actioncodec2_kernels/module.cpp

@@ -1,13 +1,25 @@
-#include <pybind11/numpy.h>
+// pybind11 entry point for the ActionCodec2 native kernels.
+//
+// Every function here has an exactly equivalent Python implementation in the
+// package.  The native module is a pure accelerator: if it is not built, the
+// package still works, and the test suite asserts the two agree bit for bit.
+
 #include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <pybind11/stl.h>
 
 #include <cstdint>
-#include <stdexcept>
-
-#include "parallel.hpp"
+#include <vector>
 
 namespace py = pybind11;
 
+void register_resample(py::module_&);
+void register_setbpe(py::module_&);
+void register_inference(py::module_&);
+void register_preprocessing(py::module_&);
+void register_grammar(py::module_&);
+void register_sparse(py::module_&);
+
 py::array_t<int32_t> so3_second_order_encode_batch(
     const py::array_t<double, py::array::c_style | py::array::forcecast>&,
     const py::array_t<int32_t, py::array::c_style | py::array::forcecast>&,
@@ -22,6 +34,15 @@ py::array_t<int32_t> additive_second_order_encode_batch(
     const py::array_t<double, py::array::c_style | py::array::forcecast>&,
     const py::array_t<float, py::array::c_style | py::array::forcecast>&, bool, int);
 
+py::tuple train_spatial_native(py::array_t<int32_t, py::array::c_style | py::array::forcecast>,
+                               py::array_t<int64_t, py::array::c_style | py::array::forcecast>,
+                               const std::vector<int32_t>&, int64_t, int);
+
+py::tuple train_temporal_native(py::iterable, int64_t,
+                                py::array_t<uint64_t, py::array::c_style | py::array::forcecast>,
+                                py::array_t<int32_t, py::array::c_style | py::array::forcecast>,
+                                int32_t, int32_t, int64_t, int64_t, int);
+
 py::array_t<float> physical_decode_batch(
     const py::array_t<float>&, const py::array_t<int32_t>&, const py::array_t<uint8_t>&,
     const py::array_t<float>&, const py::array_t<float>&,
@@ -41,30 +62,24 @@ py::array_t<int32_t> so3_encode_batch(
     const py::array_t<double, py::array::c_style | py::array::forcecast>&,
     const py::array_t<int32_t, py::array::c_style | py::array::forcecast>&,
     const py::array_t<float, py::array::c_style | py::array::forcecast>&, int32_t, bool, bool,
-    double, bool, int);
+    double, int32_t, int);
 
-void register_setbpe(py::module_& module);
-void register_resample(py::module_& module);
-void register_reconstruction(py::module_& module);
-void register_grammar(py::module_& module);
+py::tuple train_bpe_native(py::iterable, const std::vector<int32_t>&, int64_t, int64_t, int);
 
-PYBIND11_MODULE(_actioncodec2_native, module) {
-  module.doc() = "ActionCodec2 encode/decode kernels: physical stage, resampling, Set-BPE";
-  register_setbpe(module);
+py::tuple spatial_encode_instances(
+    const py::array_t<int32_t, py::array::c_style | py::array::forcecast>&,
+    const py::array_t<int64_t, py::array::c_style | py::array::forcecast>&,
+    const std::vector<int32_t>&,
+    const py::array_t<int32_t, py::array::c_style | py::array::forcecast>&, int);
+
+PYBIND11_MODULE(_native, module) {
   register_resample(module);
-  register_reconstruction(module);
+  register_setbpe(module);
+  register_inference(module);
+  register_preprocessing(module);
   register_grammar(module);
-  module.def(
-      "set_thread_budget",
-      [](int threads) {
-        if (threads < 0) throw std::invalid_argument("thread budget must be non-negative");
-        ac2::thread_budget().store(threads);
-      },
-      py::arg("threads"),
-      "Cap threads per call for this process; 0 uses every CPU it may run on.");
-  module.def("thread_budget", [] { return ac2::thread_budget().load(); });
-  module.def("plan_threads", &ac2::plan_threads, py::arg("count"), py::arg("requested"),
-             py::arg("seconds_each"), "Threads a call would use for ``count`` items.");
+  register_sparse(module);
+  module.doc() = "ActionCodec2 native kernels (Set-BPE training and batched physical stage)";
   module.def("so3_second_order_encode_batch", &so3_second_order_encode_batch,
              py::arg("table"), py::arg("bins"), py::arg("quaternions"),
              py::arg("previous_first_order"), py::arg("input_order"), py::arg("left"),
@@ -73,16 +88,39 @@ PYBIND11_MODULE(_actioncodec2_native, module) {
              py::arg("table"), py::arg("bins"), py::arg("positions"),
              py::arg("previous_first_order"), py::arg("epsilon"), py::arg("smooth"),
              py::arg("threads") = 0);
+
+  module.def("train_spatial", &train_spatial_native, py::arg("codes"), py::arg("weights"),
+             py::arg("bins"), py::arg("vocab_size"), py::arg("threads") = 0,
+             "Exact greedy spatial Set-BPE; returns (merges, report).");
+
+  module.def("train_temporal", &train_temporal_native, py::arg("instance_tokens"),
+             py::arg("total_anchors"), py::arg("spatial_support"),
+             py::arg("spatial_anchor"), py::arg("dimension"), py::arg("lookahead"),
+             py::arg("vocab_size"), py::arg("action_count"), py::arg("threads") = 0,
+             "Exact greedy rooted temporal Set-BPE; returns (merges, report).");
+
+  module.def("train_bpe", &train_bpe_native, py::arg("cell_chunks"), py::arg("bins"),
+             py::arg("vocab_size"), py::arg("max_token_cells"), py::arg("threads") = 0,
+             "Exact greedy BPE over the row-major flattened cell stream; returns (merges, report).");
+
   module.def("physical_decode_batch", &physical_decode_batch, py::arg("table"),
              py::arg("bins"), py::arg("binary"), py::arg("lower"), py::arg("upper"),
              py::arg("cells"), py::arg("initial"), py::arg("initial_velocity"),
              py::arg("order"), py::arg("threads") = 0);
+
   module.def("physical_encode_batch", &physical_encode_batch, py::arg("table"), py::arg("bins"),
              py::arg("binary"), py::arg("lower"), py::arg("upper"), py::arg("positions"),
              py::arg("initial"), py::arg("initial_velocity"), py::arg("order"),
              py::arg("threads") = 0, py::arg("epsilon") = py::array_t<float>(),
              py::arg("smooth") = false);
+
   module.def("so3_encode_batch", &so3_encode_batch, py::arg("table"), py::arg("bins"),
              py::arg("actions"), py::arg("input_order"), py::arg("rotation_vector"),
-             py::arg("left"), py::arg("limit"), py::arg("smooth"), py::arg("threads") = 0);
+             py::arg("left"), py::arg("limit"), py::arg("selection"), py::arg("threads") = 0,
+             "Batched exact SO(3) error-feedback encoding; selection 0 nearest, 1 smooth, "
+             "2 momentum_lazy.");
+
+  module.def("spatial_encode_instances", &spatial_encode_instances, py::arg("codes"),
+             py::arg("row_anchors"), py::arg("bins"), py::arg("merges"),
+             py::arg("threads") = 0);
 }

actioncodec2_kernels/parallel.hpp

@@ -16,6 +16,7 @@
 
 namespace ac2 {
 
+// Process-wide cap on threads per call; 0 means every CPU the process may use.
 inline std::atomic<int>& thread_budget() {
   static std::atomic<int> value{0};
   return value;
@@ -32,8 +33,11 @@ inline int available_cpus() {
   return std::max(1, static_cast<int>(std::thread::hardware_concurrency()));
 }
 
+// Measured cost of starting and joining one thread (seconds).
 constexpr double kThreadStartSeconds = 30e-6;
 
+// Threads for ``count`` iterations of ``seconds_each``: at most ``requested``
+// when positive, at most the process budget, and at most ``sqrt(W / s)``.
 inline int64_t plan_threads(int64_t count, int requested, double seconds_each) {
   const int budget = thread_budget().load();
   int64_t threads = budget > 0 ? budget : available_cpus();
@@ -85,4 +89,4 @@ void parallel_for(int64_t count, int threads, Body&& body) {
   if (failure) std::rethrow_exception(failure);
 }
 
-}
+}  // namespace ac2

actioncodec2_kernels/physical.cpp

@@ -1,3 +1,4 @@
+
 #include <pybind11/pybind11.h>
 #include <pybind11/numpy.h>
 #include <pybind11/stl.h>
@@ -6,7 +7,6 @@
 #include <cfloat>
 #include <cmath>
 #include <cstdint>
-#include <limits>
 #include <vector>
 
 #include "parallel.hpp"
@@ -33,7 +33,7 @@ inline int32_t pick_second_order_axis(const SecondOrderAxis& axis, double reques
     const int32_t right = std::min(insertion, axis.count - 1);
     int32_t nearest = std::fabs(request - values[left]) <= std::fabs(request - values[right])
                           ? left : right;
-
+    // Rounded distance plateaus must keep the original first-cell tie-break.
     while (nearest > 0 && std::fabs(request - values[nearest - 1]) == std::fabs(request - values[nearest]))
       --nearest;
     if constexpr (!Smooth) return nearest;
@@ -43,7 +43,7 @@ inline int32_t pick_second_order_axis(const SecondOrderAxis& axis, double reques
     const double sign = primitive > 0 ? 1 : primitive < 0 ? -1 : 0;
     return sign * direction >= 0 && std::fabs(request - primitive) <= epsilon ? preferred : nearest;
   }
-
+  // Small tables are faster to scan; arbitrary stored order also uses this path.
   int32_t nearest = 0, feasible = -1, preferred = -1;
   double nearest_distance = INFINITY, feasible_distance = INFINITY, preferred_distance = INFINITY;
   for (int32_t k = 0; k < axis.count; ++k) {
@@ -62,7 +62,7 @@ inline int32_t pick_second_order_axis(const SecondOrderAxis& axis, double reques
   return preferred >= 0 ? preferred : feasible >= 0 ? feasible : nearest;
 }
 
-}
+}  // namespace
 
 py::array_t<int32_t> additive_second_order_encode_batch(
     const py::array_t<float, py::array::c_style | py::array::forcecast>& table,
@@ -124,16 +124,17 @@ namespace {
 
 struct Layout {
   int32_t dimension;
-  int32_t width;
-  const float* table;
-  const int32_t* bins;
-  const uint8_t* binary;
+  int32_t width;                  // padded primitive-table stride
+  const float* table;             // [dimension * width], +inf padding
+  const int32_t* bins;            // [dimension]
+  const uint8_t* binary;          // [dimension], 1 for absolute two-valued axes
   const float* lower;
   const float* upper;
-  const float* epsilon;
-  bool smooth;
+  const float* epsilon;           // [dimension]; null unless `smooth` is set
+  bool smooth;                    // pick by FirstOrderCodec._select_smooth
 };
 
+// Matches PhysicalCodecBase.SOURCE_TOLERANCE.
 constexpr float kSourceTolerance = 8.0F * FLT_EPSILON;
 
 inline int32_t nearest(const Layout& layout, int32_t axis, float request) {
@@ -175,7 +176,7 @@ inline int32_t select_smooth(const Layout& layout, int32_t axis, float request,
       feasible_distance = distance;
       feasible_index = index;
     }
-
+    // np.sign(value) * direction >= 0: a zero primitive never reverses.
     const float sign = values[index] > 0.0F ? 1.0F : (values[index] < 0.0F ? -1.0F : 0.0F);
     if (!(sign * direction >= 0.0F)) continue;
     if (!any_preferred || distance < preferred_distance) {
@@ -193,7 +194,7 @@ void encode_first_order(const Layout& layout, const float* positions, int64_t st
                         const float* initial, int32_t* cells) {
   const int32_t dimension = layout.dimension;
   std::vector<float> state(initial, initial + dimension);
-
+  // Sign of the last non-zero primitive each coordinate emitted.
   std::vector<float> direction(dimension, 0.0F);
   for (int64_t step = 0; step < steps; ++step) {
     const float* target = positions + step * dimension;
@@ -297,47 +298,6 @@ Layout make_layout(const py::array_t<float>& table, const py::array_t<int32_t>&
                    const py::array_t<uint8_t>& binary, const py::array_t<float>& lower,
                    const py::array_t<float>& upper, const py::array_t<float>& epsilon,
                    bool smooth) {
-  if (table.ndim() != 2 || table.shape(0) < 1 || table.shape(1) < 1 ||
-      table.shape(0) > std::numeric_limits<int32_t>::max() ||
-      table.shape(1) > std::numeric_limits<int32_t>::max()) {
-    throw std::invalid_argument("table must have non-empty shape [D, K]");
-  }
-  for (const py::array* array : {static_cast<const py::array*>(&table),
-                               static_cast<const py::array*>(&bins),
-                               static_cast<const py::array*>(&binary),
-                               static_cast<const py::array*>(&lower),
-                               static_cast<const py::array*>(&upper),
-                               static_cast<const py::array*>(&epsilon)}) {
-    if (!(array->flags() & py::array::c_style)) {
-      throw std::invalid_argument("physical layout arrays must be C-contiguous");
-    }
-  }
-  const auto dimension = table.shape(0);
-  if (bins.ndim() != 1 || binary.ndim() != 1 || lower.ndim() != 1 || upper.ndim() != 1 ||
-      bins.size() != dimension || binary.size() != dimension ||
-      lower.size() != dimension || upper.size() != dimension) {
-    throw std::invalid_argument("layout vectors must have shape [D]");
-  }
-  if ((epsilon.size() && (epsilon.ndim() != 1 || epsilon.size() != dimension)) ||
-      (smooth && !epsilon.size())) {
-    throw std::invalid_argument("smooth selection needs epsilon with shape [D]");
-  }
-  for (py::ssize_t axis = 0; axis < dimension; ++axis) {
-    const int32_t count = bins.data()[axis];
-    if (count < 1 || count > table.shape(1) || binary.data()[axis] > 1 ||
-        !(lower.data()[axis] <= upper.data()[axis])) {
-      throw std::invalid_argument("invalid physical layout bounds or primitive count");
-    }
-    if (epsilon.size() &&
-        !(std::isfinite(epsilon.data()[axis]) && epsilon.data()[axis] >= 0)) {
-      throw std::invalid_argument("epsilon must be finite and non-negative");
-    }
-    for (int32_t cell = 0; cell < count; ++cell) {
-      if (!std::isfinite(table.data()[axis * table.shape(1) + cell])) {
-        throw std::invalid_argument("active primitives must be finite");
-      }
-    }
-  }
   Layout layout;
   layout.dimension = static_cast<int32_t>(table.shape(0));
   layout.width = static_cast<int32_t>(table.shape(1));
@@ -348,6 +308,12 @@ Layout make_layout(const py::array_t<float>& table, const py::array_t<int32_t>&
   layout.upper = upper.data();
   layout.epsilon = epsilon.size() == 0 ? nullptr : epsilon.data();
   layout.smooth = smooth && layout.epsilon != nullptr;
+  if (smooth && layout.epsilon == nullptr) {
+    throw std::invalid_argument("smooth selection needs a per-coordinate epsilon");
+  }
+  if (layout.epsilon != nullptr && epsilon.size() != layout.dimension) {
+    throw std::invalid_argument("epsilon must have one entry per coordinate");
+  }
   return layout;
 }
 
@@ -355,7 +321,7 @@ using Array2D = py::array_t<float, py::array::c_style | py::array::forcecast>;
 using Array3D = py::array_t<float, py::array::c_style | py::array::forcecast>;
 using Cells3D = py::array_t<int32_t, py::array::c_style | py::array::forcecast>;
 
-}
+}  // namespace
 
 py::array_t<float> physical_decode_batch(const py::array_t<float>& table,
                                          const py::array_t<int32_t>& bins,
@@ -365,29 +331,21 @@ py::array_t<float> physical_decode_batch(const py::array_t<float>& table,
                                          const Array2D& initial,
                                          const Array2D& initial_velocity, int32_t order,
                                          int threads) {
-  const Layout layout =
-      make_layout(table, bins, binary, lower, upper, py::array_t<float>(), false);
-  if (cells.ndim() != 3 || cells.shape(1) < 1 || cells.shape(2) != layout.dimension || initial.ndim() != 2 ||
+  if (cells.ndim() != 3 || cells.shape(2) != table.shape(0) || initial.ndim() != 2 ||
       initial_velocity.ndim() != 2 || initial.shape(0) != cells.shape(0) ||
       initial_velocity.shape(0) != cells.shape(0) || initial.shape(1) != cells.shape(2) ||
       initial_velocity.shape(1) != cells.shape(2) || (order != 1 && order != 2)) {
     throw std::invalid_argument("invalid physical decode batch shapes or order");
   }
+  const Layout layout =
+      make_layout(table, bins, binary, lower, upper, py::array_t<float>(), false);
   const int64_t batch = cells.shape(0);
   const int64_t steps = cells.shape(1);
   const int64_t dimension = layout.dimension;
-  for (py::ssize_t index = 0; index < initial.size(); ++index) {
-    if (!std::isfinite(initial.data()[index]) || !std::isfinite(initial_velocity.data()[index])) {
-      throw std::invalid_argument("initial state and velocity must be finite");
-    }
-  }
-  for (py::ssize_t index = 0; index < cells.size(); ++index) {
-    const int32_t cell = cells.data()[index];
-    if (cell < 0 || cell >= layout.bins[index % dimension]) {
-      throw std::invalid_argument("cell index is outside its primitive table");
-    }
-  }
   py::array_t<float> positions({batch, steps, dimension});
+#ifndef _OPENMP
+  (void)threads;
+#endif
   {
     py::gil_scoped_release release;
     ac2::parallel_for(batch, threads, [&](int64_t index) {
@@ -404,6 +362,7 @@ py::array_t<float> physical_decode_batch(const py::array_t<float>& table,
   return positions;
 }
 
+// One thread per chunk.  ``positions`` is [B, T, D] and ``initial`` is [B, D].
 py::array_t<int32_t> physical_encode_batch(const py::array_t<float>& table,
                                            const py::array_t<int32_t>& bins,
                                            const py::array_t<uint8_t>& binary,
@@ -414,24 +373,6 @@ py::array_t<int32_t> physical_encode_batch(const py::array_t<float>& table,
                                            int threads, const py::array_t<float>& epsilon,
                                            bool smooth) {
   const Layout layout = make_layout(table, bins, binary, lower, upper, epsilon, smooth);
-  if (positions.ndim() != 3 || positions.shape(1) < 1 ||
-      positions.shape(2) != layout.dimension || initial.ndim() != 2 ||
-      initial_velocity.ndim() != 2 || initial.shape(0) != positions.shape(0) ||
-      initial_velocity.shape(0) != positions.shape(0) ||
-      initial.shape(1) != layout.dimension || initial_velocity.shape(1) != layout.dimension ||
-      (order != 1 && order != 2)) {
-    throw std::invalid_argument("invalid physical encode batch shapes or order");
-  }
-  for (py::ssize_t index = 0; index < positions.size(); ++index) {
-    if (!std::isfinite(positions.data()[index])) {
-      throw std::invalid_argument("positions must be finite");
-    }
-  }
-  for (py::ssize_t index = 0; index < initial.size(); ++index) {
-    if (!std::isfinite(initial.data()[index]) || !std::isfinite(initial_velocity.data()[index])) {
-      throw std::invalid_argument("initial state and velocity must be finite");
-    }
-  }
   const int64_t batch = positions.shape(0);
   const int64_t steps = positions.shape(1);
   const int64_t dimension = layout.dimension;
@@ -440,6 +381,9 @@ py::array_t<int32_t> physical_encode_batch(const py::array_t<float>& table,
   const float* starts = initial.data();
   const float* velocities = initial_velocity.data();
   int32_t* output = cells.mutable_data();
+#ifndef _OPENMP
+  (void)threads;
+#endif
   {
     py::gil_scoped_release release;
     ac2::parallel_for(batch, threads, [&](int64_t index) {

actioncodec2_kernels/preprocessing.cpp

@@ -0,0 +1,281 @@
+// Continuous-segment preprocessing kernels. Policy and statistics stay in Python.
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <algorithm>
+#include <cmath>
+#include <map>
+#include <mutex>
+#include <numeric>
+#include <vector>
+#include "quaternion.hpp"
+
+namespace py = pybind11;
+namespace {
+using D64=py::array_t<double,py::array::c_style|py::array::forcecast>;
+using F32=py::array_t<float,py::array::c_style|py::array::forcecast>;
+
+double step_norm(const double* a,const double* b,int64_t width,bool single) {
+  if(single) {
+    float sum=0;
+    for(int64_t d=0;d<width;++d) {const float v=static_cast<float>(b[d])-static_cast<float>(a[d]);sum+=v*v;}
+    return std::sqrt(sum);
+  }
+  double sum=0;
+  for(int64_t d=0;d<width;++d){double v=b[d]-a[d];sum+=v*v;}
+  return std::sqrt(sum);
+}
+
+D64 step_norms(const D64& values,bool single) {
+  if(values.ndim()!=2)throw std::invalid_argument("expected [T,D] values");
+  const int64_t steps=std::max<int64_t>(0,values.shape(0)-1),width=values.shape(1);
+  D64 output(steps);
+  py::gil_scoped_release release;
+  for(int64_t t=0;t<steps;++t)output.mutable_data()[t]=step_norm(values.data()+t*width,values.data()+(t+1)*width,width,single);
+  return output;
+}
+
+std::pair<int64_t,int64_t> boundary_trim(const D64& values,double tau,int64_t idle,
+                                        int64_t context,bool legacy,bool single) {
+  if(values.ndim()!=2 || context<1)throw std::invalid_argument("invalid trimming input");
+  int64_t length=values.shape(0),width=values.shape(1);
+  if(length<2)return {0,length};
+  py::gil_scoped_release release;
+  int64_t first=-1,last=-1;
+  for(int64_t t=0;t<length-1;++t) {
+    if(step_norm(values.data()+t*width,values.data()+(t+1)*width,width,single)>tau) {
+      if(first<0)first=t;
+      last=t;
+    }
+  }
+  if(first<0)return legacy?std::make_pair(length,length):std::make_pair(int64_t{0},length);
+  const int64_t trailing=length-2-last;
+  int64_t start=first>=idle?(legacy?first+1:std::max<int64_t>(0,first-context+1)):0;
+  int64_t stop=trailing>=idle?last+2:length;
+  return {start,legacy?std::max(start,stop):std::max(start+1,stop)};
+}
+
+py::dict trim_deviation(const D64& values,int64_t start,int64_t stop,
+                        const std::string& profile,int64_t gripper) {
+  if(values.ndim()!=2 || start<0 || start>=stop || stop>values.shape(0))
+    throw std::invalid_argument("trim slice must retain at least one row");
+  const int64_t width=values.shape(1),length=values.shape(0);
+  std::vector<int64_t> motion;
+  for(int64_t d=0;d<width;++d)if(d!=gripper)motion.push_back(d);
+  if(profile=="eef" && motion.size()<6)throw std::invalid_argument("EEF needs six motion columns");
+  std::vector<double> maximum(motion.size(),0);
+  double translation=0,rotation=0;bool changed=false;
+  {
+    py::gil_scoped_release release;
+    for(int side=0;side<2;++side) {
+      int64_t begin=side?stop:0,end=side?length:start,anchor=side?stop-1:start;
+      const double* a=values.data()+anchor*width;
+      for(int64_t t=begin;t<end;++t) {
+        const double* v=values.data()+t*width;
+        for(size_t d=0;d<motion.size();++d)maximum[d]=std::max(maximum[d],std::abs(v[motion[d]]-a[motion[d]]));
+        if(gripper>=0)changed=changed || v[gripper]!=a[gripper];
+        if(profile=="eef") {
+          double distance=0,ra[3],rv[3],relative[3];
+          for(int d=0;d<3;++d) {
+            double delta=v[motion[d]]-a[motion[d]];distance+=delta*delta;
+            ra[d]=a[motion[d+3]];rv[d]=v[motion[d+3]];
+          }
+          translation=std::max(translation,std::sqrt(distance));
+          ac2q::to_rotvec(ac2q::multiply(ac2q::conjugate(ac2q::from_rotvec(ra)),ac2q::from_rotvec(rv)),relative);
+          rotation=std::max(rotation,std::sqrt(relative[0]*relative[0]+relative[1]*relative[1]+relative[2]*relative[2]));
+        }
+      }
+    }
+  }
+  py::dict result;
+  result["motion_max_error"]=maximum;result["translation_max_error_m"]=translation;
+  result["rotation_max_error_rad"]=rotation;result["gripper_changed"]=changed;
+  return result;
+}
+
+// Small normalized Vandermonde least-squares systems, solved in extended precision.
+std::vector<double> polynomial_weights(int window,int degree) {
+  const int p=degree+1;
+  std::vector<long double> basis(window*p),matrix(p*2*p,0);
+  for(int i=0;i<window;++i) {
+    long double x=2.L*i/(window-1)-1.L,v=1;
+    for(int d=0;d<p;++d){basis[i*p+d]=v;v*=x;}
+  }
+  for(int a=0;a<p;++a) {
+    for(int b=0;b<p;++b)for(int i=0;i<window;++i)matrix[a*2*p+b]+=basis[i*p+a]*basis[i*p+b];
+    matrix[a*2*p+p+a]=1;
+  }
+  for(int a=0;a<p;++a) {
+    int pivot=a;
+    for(int b=a+1;b<p;++b)if(std::abs(matrix[b*2*p+a])>std::abs(matrix[pivot*2*p+a]))pivot=b;
+    if(std::abs(matrix[pivot*2*p+a])<1e-25L)throw std::invalid_argument("singular polynomial fit");
+    for(int j=0;j<2*p;++j)std::swap(matrix[a*2*p+j],matrix[pivot*2*p+j]);
+    long double scale=matrix[a*2*p+a];
+    for(int j=0;j<2*p;++j)matrix[a*2*p+j]/=scale;
+    for(int b=0;b<p;++b)if(b!=a) {
+      const long double factor=matrix[b*2*p+a];
+      for(int j=0;j<2*p;++j)matrix[b*2*p+j]-=factor*matrix[a*2*p+j];
+    }
+  }
+  std::vector<double> weights(window*window);
+  for(int target=0;target<window;++target)for(int i=0;i<window;++i) {
+    long double w=0;
+    for(int a=0;a<p;++a)for(int b=0;b<p;++b)w+=basis[target*p+a]*matrix[a*2*p+p+b]*basis[i*p+b];
+    weights[target*window+i]=static_cast<double>(w);
+  }
+  return weights;
+}
+
+F32 savgol(const D64& values,int window,int degree) {
+  if(values.ndim()!=2 || window<3 || window%2!=1 || degree<0 || degree>=window || window>values.shape(0))
+    throw std::invalid_argument("invalid Savitzky-Golay shape or polynomial order");
+  std::vector<double> weights;
+  {
+    static std::mutex mutex;
+    static std::map<std::pair<int,int>,std::vector<double>> cache;
+    std::lock_guard<std::mutex> lock(mutex);
+    auto key=std::make_pair(window,degree);auto it=cache.find(key);
+    if(it==cache.end())it=cache.emplace(key,polynomial_weights(window,degree)).first;
+    weights=it->second;
+  }
+  const int64_t length=values.shape(0),width=values.shape(1);
+  F32 output({length,width});
+  py::gil_scoped_release release;
+  for(int64_t t=0;t<length;++t) {
+    int64_t begin=std::clamp<int64_t>(t-window/2,0,length-window),target=t-begin;
+    for(int64_t d=0;d<width;++d) {
+      double value=0;
+      for(int i=0;i<window;++i)value+=weights[target*window+i]*values.data()[(begin+i)*width+d];
+      output.mutable_data()[t*width+d]=static_cast<float>(value);
+    }
+  }
+  return output;
+}
+
+D64 median(const D64& values,int window,bool nearest) {
+  if(values.ndim()!=1 || window<1 || window%2!=1)throw std::invalid_argument("invalid median filter input");
+  const int64_t length=values.size();D64 output(length);
+  py::gil_scoped_release release;
+  std::vector<double> sample(window);
+  for(int64_t i=0;i<length;++i) {
+    for(int j=0;j<window;++j) {
+      int64_t k=i+j-window/2;
+      sample[j]=(k<0 || k>=length)?(nearest?values.data()[std::clamp<int64_t>(k,0,length-1)]:0):values.data()[k];
+    }
+    std::nth_element(sample.begin(),sample.begin()+window/2,sample.end());
+    output.mutable_data()[i]=sample[window/2];
+  }
+  return output;
+}
+
+F32 ease(const D64& values,const D64& limits) {
+  if(values.ndim()!=2 || limits.ndim()!=1 || limits.size()!=values.shape(1))
+    throw std::invalid_argument("invalid easing shape");
+  int64_t length=values.shape(0),width=values.shape(1);F32 output({length,width});
+  py::gil_scoped_release release;
+  for(int64_t d=0;d<width;++d) {
+    if(length==0)continue;
+    double debt=0,sum=0,limit=limits.data()[d];
+    output.mutable_data()[d]=static_cast<float>(values.data()[d]);
+    for(int64_t t=1;t<length;++t) {
+      double step=values.data()[t*width+d]-values.data()[(t-1)*width+d],candidate=step+debt;
+      if(std::abs(candidate)<=limit){step=candidate;debt=0;}
+      else {step=(candidate>0?1:-1)*limit;debt=candidate-step;}
+      sum+=step;output.mutable_data()[t*width+d]=static_cast<float>(values.data()[d]+sum);
+    }
+  }
+  return output;
+}
+
+F32 gripper(const D64& values,const std::string& operation,double first,double second) {
+  if(values.ndim()!=2 || values.shape(1)!=1)throw std::invalid_argument("gripper action must be [T,1]");
+  F32 output({values.shape(0),int64_t{1}});
+  py::gil_scoped_release release;
+  for(int64_t i=0;i<values.size();++i) {
+    float v=static_cast<float>(values.data()[i]);
+    if(!std::isfinite(v))throw std::invalid_argument("gripper fraction must be finite");
+    if(operation=="normalize") {
+      const float span=static_cast<float>(second-first);
+      output.mutable_data()[i]=std::clamp((v-static_cast<float>(first))/span,0.F,1.F);
+    } else if(operation=="binary")output.mutable_data()[i]=v>=static_cast<float>(first)?1.F:0.F;
+    else if(operation=="eleven")output.mutable_data()[i]=std::nearbyint(v*10.F)/10.F;
+    else throw std::invalid_argument("invalid gripper operation");
+  }
+  return output;
+}
+std::vector<double> max_change(const D64& values,const D64& reference,
+                               const std::vector<int64_t>& columns,bool single) {
+  if(values.ndim()!=2 || reference.ndim()!=2 || values.shape(0)!=reference.shape(0) || values.shape(1)!=reference.shape(1))
+    throw std::invalid_argument("mismatched change metric shapes");
+  std::vector<double> result(columns.size(),0);const int64_t width=values.shape(1);
+  py::gil_scoped_release release;
+  for(size_t i=0;i<columns.size();++i) {
+    const int64_t d=(columns[i]%width+width)%width;
+    for(int64_t t=0;t<values.shape(0);++t) {
+      double change=single?std::abs(static_cast<float>(values.data()[t*width+d])-static_cast<float>(reference.data()[t*width+d]))
+                          :std::abs(values.data()[t*width+d]-reference.data()[t*width+d]);
+      result[i]=std::max(result[i],change);
+    }
+  }
+  return result;
+}
+
+int64_t overspeed_steps(const D64& values,const D64& limits,bool single) {
+  if(values.ndim()!=2 || limits.size()!=values.shape(1))throw std::invalid_argument("invalid overspeed metric shapes");
+  int64_t count=0,width=values.shape(1);py::gil_scoped_release release;
+  for(int64_t t=1;t<values.shape(0);++t)for(int64_t d=0;d<width;++d) {
+    double delta=single?static_cast<float>(values.data()[t*width+d])-static_cast<float>(values.data()[(t-1)*width+d])
+                       :values.data()[t*width+d]-values.data()[(t-1)*width+d];
+    if(std::abs(delta)>limits.data()[d]){++count;break;}
+  }
+  return count;
+}
+
+F32 wrap_angles(const D64& values) {
+  std::vector<py::ssize_t> shape(values.shape(),values.shape()+values.ndim());F32 output(shape);
+  const double pi=std::acos(-1.);py::gil_scoped_release release;
+  for(py::ssize_t i=0;i<values.size();++i) {
+    double remainder=std::fmod(pi-values.data()[i],2*pi);if(remainder<0)remainder+=2*pi;
+    output.mutable_data()[i]=static_cast<float>(pi-remainder);
+  }
+  return output;
+}
+
+template<class T>
+void unwrap_column(const double* values,double* output,int64_t length,int64_t width,int64_t d) {
+  if(!length)return;
+  const T pi=static_cast<T>(std::acos(-1.)),period=static_cast<T>(2*std::acos(-1.));
+  T sum=0;output[d]=values[d];
+  for(int64_t t=1;t<length;++t) {
+    const T delta=static_cast<T>(values[t*width+d])-static_cast<T>(values[(t-1)*width+d]);
+    T reduced=std::fmod(static_cast<T>(delta+pi),period);if(reduced<0)reduced+=period;
+    reduced-=pi;if(reduced==-pi && delta>0)reduced=pi;
+    const T correction=std::abs(delta)<pi?0:reduced-delta;sum+=correction;
+    output[t*width+d]=static_cast<T>(static_cast<T>(values[t*width+d])+sum);
+  }
+}
+D64 unwrap_angles(const D64& values,bool single) {
+  if(values.ndim()!=2)throw std::invalid_argument("unwrap expects [T,D]");
+  int64_t length=values.shape(0),width=values.shape(1);D64 output({length,width});
+  py::gil_scoped_release release;
+  for(int64_t d=0;d<width;++d) {
+    if(single)unwrap_column<float>(values.data(),output.mutable_data(),length,width,d);
+    else unwrap_column<double>(values.data(),output.mutable_data(),length,width,d);
+  }
+  return output;
+}
+}  // namespace
+
+void register_preprocessing(py::module_& module) {
+  module.def("step_norms",&step_norms);
+  module.def("boundary_trim",&boundary_trim);
+  module.def("trim_deviation",&trim_deviation);
+  module.def("savgol_smooth",&savgol);
+  module.def("median_filter",&median);
+  module.def("ease_overspeed",&ease);
+  module.def("gripper_transform",&gripper);
+  module.def("max_change",&max_change);
+  module.def("overspeed_steps",&overspeed_steps);
+  module.def("wrap_angles",&wrap_angles);
+  module.def("unwrap_angles",&unwrap_angles);
+}

actioncodec2_kernels/quaternion.hpp

@@ -43,4 +43,4 @@ inline void to_euler(Q q, double* out) {
   out[1]=std::asin(std::clamp(2*(w*y-z*x),-1.,1.));
   out[2]=std::atan2(2*(w*z+x*y),1-2*(y*y+z*z));
 }
-}
+}  // namespace ac2q

actioncodec2_kernels/reconstruct.cpp

@@ -1,146 +0,0 @@
-#include <pybind11/numpy.h>
-#include <pybind11/pybind11.h>
-
-#include <algorithm>
-#include <cmath>
-#include <cstdint>
-#include <limits>
-#include <stdexcept>
-
-#include "parallel.hpp"
-#include "quaternion.hpp"
-
-namespace py = pybind11;
-namespace {
-using I64 = py::array_t<int64_t, py::array::c_style | py::array::forcecast>;
-using I32 = py::array_t<int32_t, py::array::c_style | py::array::forcecast>;
-using F64 = py::array_t<double, py::array::c_style | py::array::forcecast>;
-using F32 = py::array_t<float, py::array::c_style | py::array::forcecast>;
-
-void validate_reconstruction(const F64& table, const I32& bins, const I64& cells,
-                             const F32& previous, const F32& scale) {
-  if (table.ndim() != 2 || table.shape(0) < 1 || table.shape(1) < 1 ||
-      bins.ndim() != 1 || bins.size() != table.shape(0) ||
-      cells.ndim() != 3 || cells.shape(0) < 1 || cells.shape(1) < 1 ||
-      cells.shape(2) != table.shape(0) || previous.ndim() != 2 ||
-      previous.shape(0) != cells.shape(0) || previous.shape(1) != cells.shape(2) ||
-      scale.ndim() != 1 || scale.size() != cells.shape(2)) {
-    throw std::invalid_argument("invalid reconstruction shapes");
-  }
-  for (py::ssize_t d = 0; d < bins.size(); ++d) {
-    if (bins.data()[d] < 1 || bins.data()[d] > table.shape(1))
-      throw std::invalid_argument("primitive count exceeds table width");
-    if (!std::isfinite(scale.data()[d]) || scale.data()[d] <= 0)
-      throw std::invalid_argument("scale must be positive and finite");
-    for (int32_t k = 0; k < bins.data()[d]; ++k) {
-      if (!std::isfinite(table.data()[d * table.shape(1) + k]))
-        throw std::invalid_argument("active primitives must be finite");
-    }
-  }
-  for (py::ssize_t i = 0; i < previous.size(); ++i) {
-    if (!std::isfinite(previous.data()[i]))
-      throw std::invalid_argument("previous first-order values must be finite");
-  }
-  for (py::ssize_t i = 0; i < cells.size(); ++i) {
-    if (cells.data()[i] < 0 || cells.data()[i] >= bins.data()[i % bins.size()])
-      throw std::invalid_argument("cell index is outside its primitive table");
-  }
-}
-
-F32 so3_reconstruct(const F64& table, const I32& bins, const I64& cells,
-                    const F32& previous, const F32& scale, int input_order,
-                    int primitive_order, bool left, bool rotation_vector, double radius) {
-  validate_reconstruction(table, bins, cells, previous, scale);
-  if (table.shape(0) != 3 || (input_order != 0 && input_order != 1) ||
-      (primitive_order != 1 && primitive_order != 2) ||
-      std::isnan(radius) || radius <= 0)
-    throw std::invalid_argument("invalid SO(3) reconstruction order or radius");
-  const int64_t batch = cells.shape(0), steps = cells.shape(1);
-  F32 output({batch, steps, int64_t{3}});
-  auto* target = output.mutable_data();
-  py::gil_scoped_release release;
-  ac2::parallel_for(batch, 0, [&](int64_t b) {
-    double boundary[3];
-    for (int d = 0; d < 3; ++d)
-      boundary[d] = static_cast<double>(previous.data()[3 * b + d]) * scale.data()[d];
-    auto increment = rotation_vector ? ac2q::from_rotvec(boundary) : ac2q::from_euler(boundary);
-    ac2q::Q state{0, 0, 0, 1};
-    float accumulated[3] = {0, 0, 0};
-    auto apply = [&](ac2q::Q a, ac2q::Q v) {
-      return left ? ac2q::multiply(v, a) : ac2q::multiply(a, v);
-    };
-    for (int64_t t = 0; t < steps; ++t) {
-      double v[3], represented[3];
-      for (int d = 0; d < 3; ++d)
-        v[d] = table.data()[d * table.shape(1) + cells.data()[(b * steps + t) * 3 + d]];
-      if (primitive_order == 2) {
-        const double norm = std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]);
-        const double factor = std::min(1., radius / std::max(norm, std::numeric_limits<double>::min()));
-        for (double& x : v) x *= factor;
-        increment = apply(increment, ac2q::from_rotvec(v));
-        state = apply(state, increment);
-        auto q = input_order == 1 ? increment : state;
-        if (rotation_vector) ac2q::to_rotvec(q, represented);
-        else ac2q::to_euler(q, represented);
-      } else if (input_order == 1) {
-        if (rotation_vector) std::copy(v, v + 3, represented);
-        else {
-          for (double& x : v) x = static_cast<float>(x);
-          ac2q::to_euler(ac2q::from_rotvec(v), represented);
-        }
-      } else {
-        state = apply(state, ac2q::from_rotvec(v));
-        if (rotation_vector) ac2q::to_rotvec(state, represented);
-        else ac2q::to_euler(state, represented);
-      }
-      for (int d = 0; d < 3; ++d) {
-        float value = static_cast<float>(represented[d] / scale.data()[d]);
-        if (primitive_order == 1 && input_order == 1) {
-          const float next = accumulated[d] + value;
-          value = next - accumulated[d];
-          accumulated[d] = next;
-        }
-        target[(b * steps + t) * 3 + d] = value;
-      }
-    }
-  });
-  return output;
-}
-
-F32 additive_reconstruct(const F64& table, const I32& bins, const I64& cells,
-                         const F32& previous, const F32& scale, int input_order) {
-  validate_reconstruction(table, bins, cells, previous, scale);
-  if (input_order < 0 || input_order > 2)
-    throw std::invalid_argument("invalid additive reconstruction order");
-  const int64_t batch = cells.shape(0), steps = cells.shape(1), width = cells.shape(2);
-  F32 output({batch, steps, width});
-  auto* target = output.mutable_data();
-  py::gil_scoped_release release;
-  ac2::parallel_for(batch, 0, [&](int64_t b) {
-    for (int64_t d = 0; d < width; ++d) {
-      double change_sum = 0, position = 0;
-      const float boundary = previous.data()[b * width + d] * scale.data()[d];
-      for (int64_t t = 0; t < steps; ++t) {
-        const int64_t index = (b * steps + t) * width + d;
-        const double primitive = table.data()[d * table.shape(1) + cells.data()[index]];
-        change_sum += primitive;
-        const double velocity = change_sum + boundary;
-        position += velocity;
-        const double value = input_order == 2 ? primitive : input_order == 1 ? velocity : position;
-        target[index] = static_cast<float>(value / scale.data()[d]);
-      }
-    }
-  });
-  return output;
-}
-}
-
-void register_reconstruction(py::module_& module) {
-  module.def("so3_reconstruct", &so3_reconstruct, py::arg("table"), py::arg("bins"),
-             py::arg("cells"), py::arg("previous_first_order"), py::arg("scale"),
-             py::arg("input_order"), py::arg("primitive_order"), py::arg("left"),
-             py::arg("rotation_vector"), py::arg("radius"));
-  module.def("additive_reconstruct", &additive_reconstruct, py::arg("table"), py::arg("bins"),
-             py::arg("cells"), py::arg("previous_first_order"), py::arg("scale"),
-             py::arg("input_order"));
-}

actioncodec2_kernels/repacking.hpp

@@ -0,0 +1,437 @@
+// Optional lossless fixed-vocabulary packing for the native Set-BPE encoder.
+#pragma once
+#include <algorithm>
+#include <array>
+#include <cstdint>
+#include <functional>
+#include <stdexcept>
+#include <tuple>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+namespace ac2pack {
+using Instance = std::tuple<int64_t, int64_t, int64_t>;
+using Cover = std::array<uint64_t, 4>;
+
+struct CoverHash {
+  size_t operator()(const Cover& value) const {
+    size_t hash = 0;
+    for (uint64_t word : value) hash ^= std::hash<uint64_t>{}(word) +
+        0x9e3779b97f4a7c15ULL + (hash << 6) + (hash >> 2);
+    return hash;
+  }
+};
+
+struct Placement {
+  Cover cells{};
+  int64_t token;
+  int anchor, count;
+};
+
+struct Requirement { int offset; uint64_t mask, key; };
+using MatchCache = std::array<std::unordered_map<uint64_t, std::vector<int64_t>>, 8>;
+
+inline int first_free(const Cover& cover) {
+  for (int word = 0; word < 4; ++word)
+    if (~cover[word]) return word * 64 + __builtin_ctzll(~cover[word]);
+  return 256;
+}
+
+inline int free_count(const Cover& cover) {
+  int count = 0;
+  for (uint64_t word : cover) count += __builtin_popcountll(~word);
+  return count;
+}
+
+inline bool overlaps(const Cover& a, const Cover& b) {
+  for (int word = 0; word < 4; ++word) if (a[word] & b[word]) return true;
+  return false;
+}
+
+inline Cover unite(Cover a, const Cover& b) {
+  for (int word = 0; word < 4; ++word) a[word] |= b[word];
+  return a;
+}
+
+class Search {
+ public:
+  Search(const std::vector<int64_t>& bins, const std::vector<uint64_t>& support,
+         const std::vector<int64_t>& expansion,
+         const std::vector<std::vector<std::pair<int64_t, int64_t>>>& patterns,
+         const std::vector<int64_t>& roots)
+      : bins_(bins), width_(bins.size()), spatial_size_(support.size()),
+        support_(support), patterns_(patterns) {
+    if (width_ > 8) throw std::invalid_argument("packing search supports at most eight slots");
+    for (int64_t count : bins)
+      if (count > 255) throw std::invalid_argument("packing search supports at most 255 bins per slot");
+    keys_.resize(spatial_size_);
+    anchors_.resize(patterns.size());
+    spans_.resize(patterns.size(), 1);
+    counts_.resize(patterns.size());
+    for (int64_t id = 0; id < spatial_size_; ++id) {
+      for (int d = 0; d < width_; ++d) {
+        const int64_t value = expansion[id * width_ + d];
+        if (value >= 0) keys_[id] |= uint64_t(value + 1) << (8 * d);
+      }
+      lookup_.emplace(keys_[id], id);
+    }
+    temporal_by_root_.resize(spatial_size_);
+    for (size_t id = 0; id < patterns.size(); ++id) {
+      anchors_[id] = __builtin_ctzll(support[roots[id]]);
+      int count = 0;
+      for (auto [time, sid] : patterns[id]) {
+        count += __builtin_popcountll(support[sid]);
+        spans_[id] = std::max(spans_[id], int(time + 1));
+      }
+      max_cells_ = std::max(max_cells_, count);
+      counts_[id] = count;
+      if (id >= static_cast<size_t>(spatial_size_))
+        temporal_by_root_[lookup_.at(keys_[roots[id]])].push_back(id);
+    }
+    byte_masks_.resize(1 << width_);
+    for (int mask = 1; mask < (1 << width_); ++mask) {
+      const int d = __builtin_ctz(mask);
+      byte_masks_[mask] = byte_masks_[mask ^ (1 << d)] | (uint64_t{255} << (8 * d));
+    }
+    prepare_spatial_trie();
+    prepare_patterns();
+  }
+
+  std::vector<int64_t> apply(const int64_t* cells, int64_t steps,
+      const std::vector<int64_t>& original, bool spatial, bool temporal, int window, int beam) const {
+    if (!temporal) return spatial ? repack(cells, steps, original) : original;
+    auto result = local_repack(cells, steps, original, window, beam);
+    return spatial ? repack(cells, steps, result) : result;
+  }
+
+ private:
+  struct TrieNode { std::vector<int32_t> children; int64_t token = -1; };
+
+  int32_t trie_node(int dimension) {
+    const auto index = static_cast<int32_t>(trie_.size());
+    trie_.emplace_back();
+    if (dimension < width_) trie_.back().children.assign(bins_[dimension] + 1, -1);
+    return index;
+  }
+
+  void prepare_spatial_trie() {
+    for (int d = 0; d < width_; ++d) trie_roots_[d].assign(bins_[d] + 1, -1);
+    for (int64_t sid = 0; sid < spatial_size_; ++sid) {
+      if (lookup_.at(keys_[sid]) != sid) continue;
+      const int root = anchors_[sid], value = (keys_[sid] >> (8 * root)) & 255;
+      int32_t node = trie_roots_[root][value];
+      if (node < 0) {
+        node = trie_node(root + 1);
+        trie_roots_[root][value] = node;
+      }
+      for (int d = root + 1; d < width_; ++d) {
+        const int cell = (keys_[sid] >> (8 * d)) & 255;
+        int32_t child = trie_[node].children[cell];
+        if (child < 0) {
+          child = trie_node(d + 1);
+          trie_[node].children[cell] = child;
+        }
+        node = child;
+      }
+      trie_[node].token = sid;
+    }
+  }
+
+  void query_spatial_trie(int32_t node, int dimension, uint64_t row,
+                          std::vector<int64_t>& result) const {
+    if (node < 0) return;
+    if (dimension == width_) { result.push_back(trie_[node].token); return; }
+    query_spatial_trie(trie_[node].children[0], dimension + 1, row, result);
+    const int value = (row >> (8 * dimension)) & 255;
+    if (value) query_spatial_trie(trie_[node].children[value], dimension + 1, row, result);
+  }
+
+  void prepare_patterns() {
+    relative_.resize(patterns_.size());
+    requirements_.resize(patterns_.size());
+    for (size_t token = 0; token < patterns_.size(); ++token) {
+      if (spans_[token] * width_ > 256) continue;
+      std::vector<Requirement> rows(spans_[token], Requirement{0, 0, 0});
+      for (auto [offset, sid] : patterns_[token]) {
+        rows[offset].offset = offset;
+        rows[offset].mask |= byte_masks_[support_[sid]];
+        rows[offset].key |= keys_[sid];
+        uint64_t support = support_[sid];
+        while (support) {
+          const int d = __builtin_ctzll(support); support &= support - 1;
+          const int bit = offset * width_ + d;
+          relative_[token][bit / 64] |= uint64_t{1} << (bit % 64);
+        }
+      }
+      for (const auto& row : rows) if (row.mask) requirements_[token].push_back(row);
+    }
+  }
+
+  uint64_t row_key(const int64_t* row) const {
+    uint64_t key = 0;
+    for (int d = 0; d < width_; ++d)
+      if (row[d] >= 0) key |= uint64_t(row[d] + 1) << (8 * d);
+    return key;
+  }
+
+  std::vector<int64_t> matching_spatial(uint64_t row, int dim) const {
+    std::vector<int64_t> result;
+    const int value = (row >> (8 * dim)) & 255;
+    if (value) query_spatial_trie(trie_roots_[dim][value], dim + 1, row, result);
+    return result;
+  }
+
+  bool placement(int64_t token, int anchor, const std::vector<uint64_t>& rows,
+                 const Cover& fixed, Placement& out) const {
+    if (anchor + spans_[token] > static_cast<int>(rows.size())) return false;
+    for (const auto& req : requirements_[token])
+      if ((rows[anchor + req.offset] & req.mask) != req.key) return false;
+    out = Placement{}; out.token = token; out.anchor = anchor; out.count = counts_[token];
+    const int start = anchor * width_, word = start / 64, shift = start % 64;
+    const int words = (spans_[token] * width_ + 63) / 64;
+    for (int i = 0; i < words && word + i < 4; ++i) {
+      out.cells[word + i] |= relative_[token][i] << shift;
+      if (shift && word + i + 1 < 4)
+        out.cells[word + i + 1] |= relative_[token][i] >> (64 - shift);
+    }
+    return !overlaps(out.cells, fixed);
+  }
+
+  std::vector<Placement> candidates(
+      int root, const std::vector<uint64_t>& rows, const Cover& fixed,
+      MatchCache& cache) const {
+    std::vector<Placement> result;
+    const int time = root / width_, dim = root % width_;
+    auto& table = cache[dim];
+    auto found = table.find(rows[time]);
+    if (found == table.end())
+      found = table.emplace(rows[time], matching_spatial(rows[time], dim)).first;
+    for (int64_t sid : found->second) {
+      Placement offer;
+      if (placement(sid, time, rows, fixed, offer)) result.push_back(offer);
+      for (int64_t token : temporal_by_root_[sid])
+        if (placement(token, time, rows, fixed, offer)) result.push_back(offer);
+    }
+    std::sort(result.begin(), result.end(), [](const Placement& a, const Placement& b) {
+      return std::make_pair(-a.count, a.token) < std::make_pair(-b.count, b.token);
+    });
+    return result;
+  }
+
+  using Options = std::function<const std::vector<Placement>&(int)>;
+
+  std::vector<Instance> greedy_cover(const Options& options, Cover occupied, int bound) const {
+    std::vector<Instance> result;
+    while (first_free(occupied) != 256 && static_cast<int>(result.size()) + 1 < bound) {
+      const auto& offers = options(first_free(occupied));
+      auto choice = std::find_if(offers.begin(), offers.end(), [&](const Placement& offer) {
+        return !overlaps(occupied, offer.cells);
+      });
+      if (choice == offers.end()) return {};
+      result.emplace_back(choice->anchor, anchors_[choice->token], choice->token);
+      occupied = unite(occupied, choice->cells);
+    }
+    return first_free(occupied) == 256 ? result : std::vector<Instance>{};
+  }
+
+  std::vector<Instance> beam_cover(const Options& options,
+                                   const Cover& fixed, int bound, int beam) const {
+    if (beam == 1) return greedy_cover(options, fixed, bound);
+    struct Node { Cover covered; int parent; int64_t token; int anchor, remaining, root; };
+    std::vector<Node> arena{{fixed, -1, -1, -1, free_count(fixed), first_free(fixed)}};
+    arena.reserve(beam * bound + 1);
+    std::vector<int> frontier{0};
+    frontier.reserve(beam);
+    std::array<Node, 64> next;
+    auto better = [](const Node& a, const Node& b) {
+      if (a.remaining != b.remaining) return a.remaining < b.remaining;
+      return a.root != b.root ? a.root > b.root : a.covered < b.covered;
+    };
+    for (int depth = 0; depth + 1 < bound && !frontier.empty(); ++depth) {
+      int keep = 0;
+      for (int index : frontier) {
+        for (const auto& offer : options(arena[index].root)) {
+          if (overlaps(arena[index].covered, offer.cells)) continue;
+          Cover covered = unite(arena[index].covered, offer.cells);
+          const int remaining = arena[index].remaining - offer.count;
+          if (!remaining) {
+            std::vector<Instance> result{{offer.anchor, anchors_[offer.token], offer.token}};
+            for (int p = index; arena[p].parent >= 0; p = arena[p].parent)
+              result.emplace_back(arena[p].anchor, anchors_[arena[p].token], arena[p].token);
+            std::reverse(result.begin(), result.end());
+            return result;
+          }
+          if (keep == beam && remaining > next[keep - 1].remaining) continue;
+          Node node{covered, index, offer.token, offer.anchor, remaining, first_free(covered)};
+          if (keep == beam && !better(node, next[keep - 1])) continue;
+          bool duplicate = false;
+          for (int i = 0; i < keep; ++i) if (covered == next[i].covered) { duplicate = true; break; }
+          if (duplicate) continue;
+          int position = keep < beam ? keep++ : keep - 1;
+          while (position && better(node, next[position - 1])) {
+            next[position] = next[position - 1]; --position;
+          }
+          next[position] = node;
+        }
+      }
+      frontier.clear();
+      for (int i = 0; i < keep; ++i) {
+        frontier.push_back(arena.size());
+        arena.push_back(std::move(next[i]));
+      }
+    }
+    return {};
+  }
+
+  std::vector<int64_t> local_repack(const int64_t* cells, int64_t steps,
+                                   const std::vector<int64_t>& stream, int window, int beam) const {
+    uint64_t full = 0;
+    for (int d = 0; d < width_; ++d) if (cells[d] >= 0) full |= uint64_t{1} << d;
+    std::vector<uint64_t> coverage(steps, 0), fixed(steps, 0);
+    std::vector<std::vector<Instance>> blocks((steps + window - 1) / window);
+    std::vector<Instance> result;
+    int64_t front = 0;
+    for (int64_t token : stream) {
+      while (front < steps && coverage[front] == full) ++front;
+      const bool crossing = front / window != (front + spans_[token] - 1) / window;
+      Instance instance{front, anchors_[token], token};
+      if (crossing) result.push_back(instance);
+      else blocks[front / window].push_back(instance);
+      for (auto [offset, sid] : patterns_[token]) {
+        coverage[front + offset] |= support_[sid];
+        if (crossing) fixed[front + offset] |= support_[sid];
+      }
+    }
+    MatchCache cache;
+    for (size_t block = 0; block < blocks.size(); ++block) {
+      const int start = block * window, length = std::min<int64_t>(window, steps - start);
+      Cover occupied; occupied.fill(~uint64_t{0});
+      std::vector<uint64_t> rows(length);
+      for (int t = 0; t < length; ++t) {
+        rows[t] = row_key(cells + (start + t) * width_);
+        for (int d = 0; d < width_; ++d) {
+          const int bit = t * width_ + d;
+          if ((full & (uint64_t{1} << d)) && !(fixed[start + t] & (uint64_t{1} << d)))
+            occupied[bit / 64] &= ~(uint64_t{1} << (bit % 64));
+        }
+      }
+      std::vector<Instance> improved;
+      const int bound = blocks[block].size();
+      if ((free_count(occupied) + max_cells_ - 1) / max_cells_ < bound) {
+        std::vector<std::vector<Placement>> options(length * width_);
+        std::vector<char> ready(length * width_, 0);
+        auto get_options = [&](int root) -> const std::vector<Placement>& {
+          if (!ready[root]) {
+            options[root] = candidates(root, rows, occupied, cache);
+            ready[root] = 1;
+          }
+          return options[root];
+        };
+        improved = beam_cover(get_options, occupied, bound, beam);
+      }
+      if (improved.empty()) result.insert(result.end(), blocks[block].begin(), blocks[block].end());
+      else for (auto [time, root, token] : improved) result.emplace_back(start + time, root, token);
+      for (auto& table : cache) if (table.size() > 4096) table.clear();
+    }
+    std::sort(result.begin(), result.end());
+    std::vector<int64_t> tokens;
+    for (auto [time, root, token] : result) tokens.push_back(token);
+    if (tokens.size() > stream.size()) throw std::runtime_error("local packing regressed");
+    return tokens;
+  }
+
+  std::vector<int64_t> optimal(uint64_t row, int mask,
+                               const std::vector<int64_t>& incumbent) const {
+    if (incumbent.size() < 2) return incumbent;
+    std::array<int64_t, 256> matched;
+    matched.fill(-2);
+    auto match = [&](int subset) {
+      int64_t& token = matched[subset];
+      if (token == -2) {
+        const auto found = lookup_.find(row & byte_masks_[subset]);
+        token = found == lookup_.end() ? -1 : found->second;
+      }
+      return token;
+    };
+    if (int64_t token = match(mask); token >= 0) return {token};
+    if (incumbent.size() == 2) return incumbent;
+    const int bit = mask & -mask;
+    for (int sub = (mask - 1) & mask; sub; sub = (sub - 1) & mask) {
+      if (!(sub & bit)) continue;
+      const int64_t left = match(sub), right = match(mask ^ sub);
+      if (left >= 0 && right >= 0) return {left, right};
+    }
+    if (incumbent.size() == 3) return incumbent;
+    std::array<int, 256> costs, choices;
+    costs.fill(99); choices.fill(0); costs[0] = 0;
+    std::function<int(int)> solve = [&](int remaining) {
+      if (costs[remaining] != 99) return costs[remaining];
+      int best = 98;
+      const int first = remaining & -remaining;
+      for (int sub = remaining; sub; sub = (sub - 1) & remaining) {
+        if (!(sub & first) || match(sub) < 0) continue;
+        const int cost = 1 + solve(remaining ^ sub);
+        if (cost < best) { best = cost; choices[remaining] = sub; }
+      }
+      return costs[remaining] = best;
+    };
+    if (solve(mask) >= static_cast<int>(incumbent.size())) return incumbent;
+    std::vector<int64_t> result;
+    for (int remaining = mask; remaining; remaining ^= choices[remaining])
+      result.push_back(match(choices[remaining]));
+    return result;
+  }
+
+  std::vector<int64_t> repack(const int64_t* cells, int64_t steps,
+                            const std::vector<int64_t>& stream) const {
+    uint64_t full = 0;
+    for (int d = 0; d < width_; ++d) if (cells[d] >= 0) full |= uint64_t{1} << d;
+    std::vector<uint64_t> covered(steps, 0), fixed(steps, 0);
+    std::vector<std::vector<int64_t>> old_spatial(steps), temporal(steps);
+    int64_t front = 0;
+    for (int64_t token : stream) {
+      while (front < steps && covered[front] == full) ++front;
+      if (front == steps || token < 0 || token >= static_cast<int64_t>(patterns_.size()) ||
+          anchors_[token] != __builtin_ctzll(full & ~covered[front]))
+        throw std::invalid_argument("invalid original token stream");
+      const bool frozen = token >= spatial_size_;
+      (frozen ? temporal[front] : old_spatial[front]).push_back(token);
+      for (auto [offset, id] : patterns_[token]) {
+        if (front + offset >= steps || (covered[front + offset] & support_[id]))
+          throw std::invalid_argument("invalid original token coverage");
+        covered[front + offset] |= support_[id];
+        if (frozen) fixed[front + offset] |= support_[id];
+      }
+    }
+    std::vector<int64_t> result;
+    result.reserve(stream.size());
+    for (int64_t time = 0; time < steps; ++time) {
+      if (covered[time] != full) throw std::invalid_argument("incomplete original stream");
+      auto tokens = optimal(row_key(cells + time * width_), full & ~fixed[time], old_spatial[time]);
+      tokens.insert(tokens.end(), temporal[time].begin(), temporal[time].end());
+      std::sort(tokens.begin(), tokens.end(), [&](int64_t a, int64_t b) {
+        return std::make_pair(anchors_[a], a) < std::make_pair(anchors_[b], b);
+      });
+      result.insert(result.end(), tokens.begin(), tokens.end());
+    }
+    if (result.size() > stream.size()) throw std::runtime_error("packing regressed");
+    return result;
+  }
+
+  const std::vector<int64_t>& bins_;
+  int width_;
+  int64_t spatial_size_;
+  const std::vector<uint64_t>& support_;
+  const std::vector<std::vector<std::pair<int64_t, int64_t>>>& patterns_;
+  std::vector<uint64_t> keys_, byte_masks_;
+  std::vector<int> anchors_, spans_, counts_;
+  std::vector<Cover> relative_;
+  std::vector<std::vector<Requirement>> requirements_;
+  std::vector<std::vector<int64_t>> temporal_by_root_;
+  std::vector<TrieNode> trie_;
+  std::array<std::vector<int32_t>, 8> trie_roots_;
+  int max_cells_ = 1;
+  std::unordered_map<uint64_t, int64_t> lookup_;
+};
+}  // namespace ac2pack

actioncodec2_kernels/replay.cpp

@@ -0,0 +1,160 @@
+// Parallel replay of a learned Spatial Set-BPE vocabulary over action cells.
+
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <functional>
+#include <queue>
+#include <tuple>
+#include <vector>
+
+#ifdef _OPENMP
+#include <omp.h>
+#endif
+
+namespace py = pybind11;
+
+namespace {
+
+struct Replay {
+  int32_t dimension;
+  int32_t vocab_size;
+  std::vector<int32_t> offsets;
+  std::vector<int32_t> rule;
+  std::vector<int32_t> anchor;
+
+  int32_t child(int32_t left, int32_t right) const {
+    if (left > right) std::swap(left, right);
+    return rule[static_cast<size_t>(left) * vocab_size + right];
+  }
+
+  std::vector<int32_t> encode(const int32_t* code) const {
+    std::vector<int32_t> active;
+    active.reserve(dimension);
+    for (int32_t axis = 0; axis < dimension; ++axis) {
+      if (code[axis] < 0) continue;
+      active.push_back(offsets[axis] + code[axis]);
+    }
+    using Candidate = std::tuple<int32_t, int32_t, int32_t>;
+    std::priority_queue<Candidate, std::vector<Candidate>, std::greater<Candidate>> heap;
+    auto offer = [&](int32_t left, int32_t right) {
+      const int32_t value = child(left, right);
+      if (value >= 0) heap.emplace(value, std::min(left, right), std::max(left, right));
+    };
+    for (int32_t a = 0; a < static_cast<int32_t>(active.size()); ++a) {
+      for (int32_t b = a + 1; b < static_cast<int32_t>(active.size()); ++b) {
+        offer(active[a], active[b]);
+      }
+    }
+    while (!heap.empty()) {
+      const auto [child_token, left, right] = heap.top();
+      heap.pop();
+      auto left_it = std::find(active.begin(), active.end(), left);
+      auto right_it = std::find(active.begin(), active.end(), right);
+      if (left_it == active.end() || right_it == active.end()) continue;
+      *left_it = child_token;
+      active.erase(right_it);
+      for (const int32_t other : active) {
+        if (other != child_token) offer(child_token, other);
+      }
+    }
+    std::sort(active.begin(), active.end(), [&](int32_t left, int32_t right) {
+      if (anchor[left] != anchor[right]) return anchor[left] < anchor[right];
+      return left < right;
+    });
+    return active;
+  }
+};
+
+}  // namespace
+
+py::tuple spatial_encode_instances(
+    const py::array_t<int32_t, py::array::c_style | py::array::forcecast>& codes,
+    const py::array_t<int64_t, py::array::c_style | py::array::forcecast>& row_anchors,
+    const std::vector<int32_t>& bins,
+    const py::array_t<int32_t, py::array::c_style | py::array::forcecast>& merges,
+    int threads) {
+  if (codes.ndim() != 2 || row_anchors.ndim() != 1 ||
+      row_anchors.shape(0) != codes.shape(0) ||
+      codes.shape(1) != static_cast<int64_t>(bins.size()) || merges.ndim() != 2 ||
+      merges.shape(1) != 2) {
+    throw std::invalid_argument("invalid spatial replay inputs");
+  }
+  Replay replay;
+  replay.dimension = static_cast<int32_t>(bins.size());
+  replay.offsets.assign(replay.dimension + 1, 0);
+  for (int32_t axis = 0; axis < replay.dimension; ++axis) {
+    replay.offsets[axis + 1] = replay.offsets[axis] + bins[axis];
+  }
+  const int32_t atoms = replay.offsets.back();
+  replay.vocab_size = atoms + static_cast<int32_t>(merges.shape(0));
+  replay.rule.assign(static_cast<size_t>(replay.vocab_size) * replay.vocab_size, -1);
+  replay.anchor.assign(replay.vocab_size, -1);
+  for (int32_t axis = 0; axis < replay.dimension; ++axis) {
+    for (int32_t token = replay.offsets[axis]; token < replay.offsets[axis + 1]; ++token) {
+      replay.anchor[token] = axis;
+    }
+  }
+  auto rules = merges.unchecked<2>();
+  for (int32_t index = 0; index < merges.shape(0); ++index) {
+    const int32_t left = rules(index, 0), right = rules(index, 1), child = atoms + index;
+    if (left < 0 || right < 0 || left >= child || right >= child) {
+      throw std::invalid_argument("invalid spatial merge rule");
+    }
+    replay.rule[static_cast<size_t>(std::min(left, right)) * replay.vocab_size +
+                std::max(left, right)] = child;
+    replay.anchor[child] = std::min(replay.anchor[left], replay.anchor[right]);
+  }
+
+  const int64_t rows = codes.shape(0);
+  std::vector<int32_t> counts(rows);
+  const int32_t* source = codes.data();
+  for (int64_t row = 0; row < rows; ++row) {
+    int32_t active = 0;
+    for (int32_t axis = 0; axis < replay.dimension; ++axis) {
+      const int32_t value = source[row * replay.dimension + axis];
+      if (value < -1 || value >= bins[axis]) {
+        throw std::invalid_argument("a cell index is outside its coordinate's bin count");
+      }
+      active += value >= 0;
+    }
+    if (active == 0) throw std::invalid_argument("every row needs an active coordinate");
+  }
+#ifndef _OPENMP
+  (void)threads;
+#endif
+  {
+    py::gil_scoped_release release;
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static) num_threads(threads > 0 ? threads : omp_get_max_threads())
+#endif
+    for (int64_t row = 0; row < rows; ++row) {
+      counts[row] = static_cast<int32_t>(
+          replay.encode(source + row * replay.dimension).size());
+    }
+  }
+  std::vector<int64_t> offsets(rows + 1, 0);
+  for (int64_t row = 0; row < rows; ++row) offsets[row + 1] = offsets[row] + counts[row];
+  py::array_t<int32_t> tokens(offsets.back());
+  py::array_t<int64_t> anchors(offsets.back());
+  int32_t* output_tokens = tokens.mutable_data();
+  int64_t* output_anchors = anchors.mutable_data();
+  const int64_t* input_anchors = row_anchors.data();
+  {
+    py::gil_scoped_release release;
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static) num_threads(threads > 0 ? threads : omp_get_max_threads())
+#endif
+    for (int64_t row = 0; row < rows; ++row) {
+      const auto encoded = replay.encode(source + row * replay.dimension);
+      const int64_t begin = offsets[row];
+      for (int32_t index = 0; index < static_cast<int32_t>(encoded.size()); ++index) {
+        output_tokens[begin + index] = encoded[index];
+        output_anchors[begin + index] = input_anchors[row];
+      }
+    }
+  }
+  return py::make_tuple(tokens, anchors);
+}

actioncodec2_kernels/resample.cpp

@@ -1,3 +1,4 @@
+
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
 
@@ -21,6 +22,8 @@ struct Quat {
   double x, y, z, w;
 };
 
+// ---- runtime ``_so3`` quaternion math (xyzw), NumPy evaluation order ----
+
 Quat so3_normalize(const Quat& q) {
   const double norm = std::sqrt(q.x * q.x + q.y * q.y + q.z * q.z + q.w * q.w);
   if (norm < 1e-12) throw std::invalid_argument("a pose contains a zero quaternion");
@@ -53,6 +56,8 @@ void so3_to_rotvec(const Quat& raw, double* out) {
   out[2] = q.z * scale;
 }
 
+// ---- scipy ``Rotation`` operations used by RotationSpline ----
+
 Quat sp_from_quat(const Quat& q) {
   const double norm = std::sqrt(q.x * q.x + q.y * q.y + q.z * q.z + q.w * q.w);
   if (!(norm > 0.0)) throw std::invalid_argument("found zero norm quaternions in `quat`");
@@ -97,6 +102,9 @@ Quat sp_from_rotvec(const double* v) {
   return {scale * v[0], scale * v[1], scale * v[2], std::cos(angle / 2.0)};
 }
 
+// ---- piecewise cubics (scipy PPoly with CubicHermiteSpline coefficients) ----
+
+// PPoly interval: x[i] <= t < x[i+1], extrapolating with the end intervals.
 int64_t ppoly_interval(const Vec& x, double t) {
   const int64_t n = static_cast<int64_t>(x.size());
   if (t >= x[n - 1]) return n - 2;
@@ -104,6 +112,8 @@ int64_t ppoly_interval(const Vec& x, double t) {
   return static_cast<int64_t>(std::upper_bound(x.begin(), x.end(), t) - x.begin()) - 1;
 }
 
+// scipy ``evaluate_poly1`` for a cubic: sum of c[k-1-p] * s**p, powers built by
+// repeated multiplication.
 double ppoly_value(const double* c, double s) {
   double result = 0.0, z = 1.0;
   result = result + c[3] * z;
@@ -116,6 +126,8 @@ double ppoly_value(const double* c, double s) {
   return result;
 }
 
+// Hermite coefficients per interval and column: c[(k*4 + p)*cols + a] with
+// p = 0..3 the descending powers, as CubicHermiteSpline stacks them.
 Vec hermite_coefficients(const Vec& x, const Vec& y, const Vec& slopes, int64_t cols) {
   const int64_t n = static_cast<int64_t>(x.size());
   Vec c(static_cast<size_t>((n - 1) * 4 * cols));
@@ -148,6 +160,7 @@ void ppoly_evaluate(const Vec& x, const Vec& c, int64_t cols, const Vec& times,
   }
 }
 
+// LAPACK dgtsv (reference algorithm) for several right-hand sides b[n][cols].
 void gtsv(Vec dl, Vec d, Vec du, Vec& b, int64_t cols) {
   const int64_t n = static_cast<int64_t>(d.size());
   auto pivot = [&](int64_t i, bool last) {
@@ -188,6 +201,7 @@ void gtsv(Vec dl, Vec d, Vec du, Vec& b, int64_t cols) {
   }
 }
 
+// scipy CubicSpline(bc_type="natural") slopes at the knots.
 Vec natural_cubic_slopes(const Vec& x, const Vec& y, int64_t cols) {
   const int64_t n = static_cast<int64_t>(x.size());
   Vec dx(n - 1), slope((n - 1) * cols);
@@ -195,7 +209,7 @@ Vec natural_cubic_slopes(const Vec& x, const Vec& y, int64_t cols) {
     dx[k] = x[k + 1] - x[k];
     for (int64_t a = 0; a < cols; ++a) slope[k * cols + a] = (y[(k + 1) * cols + a] - y[k * cols + a]) / dx[k];
   }
-
+  // Banded A (3 x n): A[0] superdiagonal, A[1] diagonal, A[2] subdiagonal.
   Vec upper(n, 0.0), diagonal(n, 0.0), lower(n, 0.0), b(n * cols);
   for (int64_t k = 1; k < n - 1; ++k) diagonal[k] = 2.0 * (dx[k - 1] + dx[k]);
   for (int64_t k = 2; k < n; ++k) upper[k] = dx[k - 2];
@@ -204,7 +218,7 @@ Vec natural_cubic_slopes(const Vec& x, const Vec& y, int64_t cols) {
     for (int64_t a = 0; a < cols; ++a)
       b[k * cols + a] = 3.0 * (dx[k] * slope[(k - 1) * cols + a] + dx[k - 1] * slope[k * cols + a]);
   }
-
+  // Natural ends: second derivative zero (bc value 0.0).
   diagonal[0] = 2.0 * dx[0];
   upper[1] = dx[0];
   diagonal[n - 1] = 2.0 * dx[n - 2];
@@ -230,6 +244,7 @@ double pchip_edge(double h0, double h1, double m0, double m1) {
   return d;
 }
 
+// scipy PchipInterpolator derivatives at the knots.
 Vec pchip_slopes(const Vec& x, const Vec& y, int64_t cols) {
   const int64_t n = static_cast<int64_t>(x.size());
   Vec h(n - 1), m((n - 1) * cols), dk(n * cols, 0.0);
@@ -262,6 +277,8 @@ Vec pchip_slopes(const Vec& x, const Vec& y, int64_t cols) {
   return dk;
 }
 
+// ---- scipy RotationSpline ----
+
 using Mat3 = std::array<double, 9>;
 
 Mat3 skew(const double* v) { return {0.0, -v[2], v[1], v[2], 0.0, -v[0], -v[1], v[0], 0.0}; }
@@ -363,7 +380,7 @@ void gbsv(int64_t n, int64_t kl, int64_t ku, Vec& ab, Vec& b) {
       }
     }
   }
-
+  // Forward: L, with the row interchanges.
   for (int64_t j = 0; j + 1 < n; ++j) {
     const int64_t lm = std::min(kl, n - 1 - j);
     if (pivots[j] != j) std::swap(b[pivots[j]], b[j]);
@@ -372,7 +389,7 @@ void gbsv(int64_t n, int64_t kl, int64_t ku, Vec& ab, Vec& b) {
       for (int64_t i = 1; i <= lm; ++i) b[j + i] = b[j + i] + at(kv + i, j) * temp;
     }
   }
-
+  // Backward: U, a band of width kl+ku above the diagonal (dtbsv).
   const int64_t k = kl + ku;
   for (int64_t j = n - 1; j >= 0; --j) {
     if (b[j] != 0.0) {
@@ -386,7 +403,7 @@ void gbsv(int64_t n, int64_t kl, int64_t ku, Vec& ab, Vec& b) {
 struct RotationSpline {
   Vec times;
   std::vector<Quat> rotations;
-  Vec coefficients;
+  Vec coefficients;  // ppoly layout, 3 columns
 
   RotationSpline(const Vec& knots, const std::vector<Quat>& raw) : times(knots) {
     const int64_t n = static_cast<int64_t>(knots.size());
@@ -416,8 +433,8 @@ struct RotationSpline {
     }
   }
 
-
-
+  // scipy RotationSpline._solve_for_angular_rates; ``rates`` comes back as the
+  // knot angular rates, ``rotvecs_dot`` as the segment rotation-vector rates.
   static void solve_rates(const Vec& dt, Vec& rates, const Vec& rotvecs, Vec& rotvecs_dot) {
     const int64_t m = static_cast<int64_t>(dt.size());
     const double first[3] = {rates[0], rates[1], rates[2]};
@@ -426,7 +443,7 @@ struct RotationSpline {
       a[k] = rate_to_rotvec_dot(&rotvecs[k * 3]);
       a_inv[k] = rotvec_dot_to_rate(&rotvecs[k * 3]);
     }
-
+    // Block-tridiagonal system of m-1 unknown 3-vectors in (5, 5) band form.
     const int64_t size = 3 * (m - 1), kl = 5, ku = 5, ldab = 2 * kl + ku + 1;
     Vec band(ldab * size, 0.0);
     auto set = [&](int64_t i, int64_t j, double value) { band[(kl + ku + i - j) + ldab * j] = value; };
@@ -474,7 +491,7 @@ struct RotationSpline {
       if (converged) break;
     }
     for (int64_t k = 0; k < m; ++k) matvec(a[k], &rates[k * 3], &rotvecs_dot[k * 3]);
-
+    // Knot rates: the fixed first rate, then the solved interior ones.
     Vec shifted(m * 3);
     for (int r = 0; r < 3; ++r) shifted[r] = first[r];
     for (int64_t i = 0; i < 3 * (m - 1); ++i) shifted[3 + i] = rates[i];
@@ -516,6 +533,8 @@ Quat slerp(const Quat& left, const Quat& right, double weight) {
   return so3_normalize(out);
 }
 
+// ---- runtime resample_uniform ----
+
 enum Kind { kLinear = 0, kHold = 1, kPose = 2 };
 enum Method { kLinearMethod = 0, kPchip = 1, kCubic = 2 };
 
@@ -525,7 +544,8 @@ int64_t resampled_horizon(int64_t length, double source_hz, double target_hz) {
 
 int64_t output_width(int kind, int64_t width) { return kind == kPose ? 7 : width; }
 
-void resample_row(const float* input, int64_t steps, int64_t width, double source_hz, double target_hz,
+// One row: float32 [T, D] in, float32 [count, D'] out.
+void resample_row(const double* input, int64_t steps, int64_t width, double source_hz, double target_hz,
                   int kind, int64_t count, bool identity, int method, float* output) {
   Vec values(steps * width);
   for (int64_t i = 0; i < steps * width; ++i) values[i] = input[i];
@@ -565,7 +585,7 @@ void resample_row(const float* input, int64_t steps, int64_t width, double sourc
     } else if (method == kCubic) {
       const Vec c = hermite_coefficients(source_time, continuous, natural_cubic_slopes(source_time, continuous, cols), cols);
       ppoly_evaluate(source_time, c, cols, interpolation_time, out_width, 0, result.data());
-
+      // Guard: fall back to PCHIP on axes that materially leave their envelope.
       std::vector<int64_t> fallback;
       for (int64_t a = 0; a < cols; ++a) {
         double lower = continuous[a], upper = continuous[a];
@@ -635,14 +655,15 @@ int parse_method(const std::string& method) {
 }
 
 using F32 = py::array_t<float, py::array::c_style | py::array::forcecast>;
+using F64 = py::array_t<double, py::array::c_style | py::array::forcecast>;
 
-void check_finite(const F32& values) {
-  const float* data = values.data();
+void check_finite(const F64& values) {
+  const double* data = values.data();
   for (py::ssize_t i = 0; i < values.size(); ++i)
     if (!std::isfinite(data[i])) throw std::invalid_argument("values must be a non-empty finite [T, D] array");
 }
 
-F32 resample_uniform_batch(const F32& values, double source_hz, double target_hz, const std::string& kind_name,
+F32 resample_uniform_batch(const F64& values, double source_hz, double target_hz, const std::string& kind_name,
                            int64_t target_length, const std::string& method_name, int threads) {
   if (values.ndim() != 3 || values.shape(0) < 1 || values.shape(1) < 1)
     throw std::invalid_argument("values must have non-empty shape [B, T, D]");
@@ -658,7 +679,7 @@ F32 resample_uniform_batch(const F32& values, double source_hz, double target_hz
   if (count < 1) throw std::invalid_argument("target_length must be a positive integer");
   const int64_t out_width = output_width(kind, width);
   F32 output({batch, count, out_width});
-  const float* in = values.data();
+  const double* in = values.data();
   float* out = output.mutable_data();
   {
     py::gil_scoped_release release;
@@ -670,14 +691,15 @@ F32 resample_uniform_batch(const F32& values, double source_hz, double target_hz
   return output;
 }
 
+// Elementwise pose maps over [B, T, D] float32 arrays.
 template <class Row>
-F32 map_rows(const F32& values, int64_t in_width, int64_t out_width, int threads, Row&& row) {
+F32 map_rows(const F64& values, int64_t in_width, int64_t out_width, int threads, Row&& row) {
   if (values.ndim() != 3 || values.shape(2) != in_width || values.shape(1) < 1)
     throw std::invalid_argument("expected shape [B, T, " + std::to_string(in_width) + "]");
   check_finite(values);
   const int64_t batch = values.shape(0), steps = values.shape(1);
   F32 output({batch, steps, out_width});
-  const float* in = values.data();
+  const double* in = values.data();
   float* out = output.mutable_data();
   {
     py::gil_scoped_release release;
@@ -688,8 +710,9 @@ F32 map_rows(const F32& values, int64_t in_width, int64_t out_width, int threads
   return output;
 }
 
-F32 integrate_eef_deltas_batch(const F32& deltas, int threads) {
-  return map_rows(deltas, 6, 7, threads, [](const float* in, int64_t steps, float* out) {
+// runtime integrate_eef_deltas: [T, 6] body increments -> [T, 7] absolute poses.
+F32 integrate_eef_deltas_batch(const F64& deltas, int threads) {
+  return map_rows(deltas, 6, 7, threads, [](const double* in, int64_t steps, float* out) {
     double position[3] = {0.0, 0.0, 0.0};
     Quat current{0.0, 0.0, 0.0, 1.0};
     for (int64_t t = 0; t < steps; ++t) {
@@ -707,8 +730,9 @@ F32 integrate_eef_deltas_batch(const F32& deltas, int threads) {
   });
 }
 
-F32 eef_pose_from_rotvec_batch(const F32& poses, int threads) {
-  return map_rows(poses, 6, 7, threads, [](const float* in, int64_t steps, float* out) {
+// runtime eef_pose_from_components(..., "rotation_vector").
+F32 eef_pose_from_rotvec_batch(const F64& poses, int threads) {
+  return map_rows(poses, 6, 7, threads, [](const double* in, int64_t steps, float* out) {
     for (int64_t t = 0; t < steps; ++t) {
       for (int a = 0; a < 3; ++a) out[t * 7 + a] = in[t * 6 + a];
       const Quat q = so3_from_rotvec(in[t * 6 + 3], in[t * 6 + 4], in[t * 6 + 5]);
@@ -720,8 +744,9 @@ F32 eef_pose_from_rotvec_batch(const F32& poses, int threads) {
   });
 }
 
-F32 eef_pose_absolute_batch(const F32& poses, int threads) {
-  return map_rows(poses, 7, 6, threads, [](const float* in, int64_t steps, float* out) {
+// runtime eef_pose_absolute: [T, 7] xyz + xyzw -> [T, 6] xyz + rotation vector.
+F32 eef_pose_absolute_batch(const F64& poses, int threads) {
+  return map_rows(poses, 7, 6, threads, [](const double* in, int64_t steps, float* out) {
     for (int64_t t = 0; t < steps; ++t) {
       for (int a = 0; a < 3; ++a) out[t * 6 + a] = in[t * 7 + a];
       double rotvec[3];
@@ -731,8 +756,9 @@ F32 eef_pose_absolute_batch(const F32& poses, int threads) {
   });
 }
 
-F32 eef_pose_deltas_batch(const F32& poses, int threads) {
-  return map_rows(poses, 7, 6, threads, [](const float* in, int64_t steps, float* out) {
+// runtime eef_pose_deltas: local [T, 7] path -> [T, 6] body increments.
+F32 eef_pose_deltas_batch(const F64& poses, int threads) {
+  return map_rows(poses, 7, 6, threads, [](const double* in, int64_t steps, float* out) {
     Quat previous{0.0, 0.0, 0.0, 1.0};
     for (int64_t t = 0; t < steps; ++t) {
       for (int a = 0; a < 3; ++a) {
@@ -748,8 +774,10 @@ F32 eef_pose_deltas_batch(const F32& poses, int threads) {
   });
 }
 
-F32 rotvec_left_to_right_batch(const F32& increments, int threads) {
-  return map_rows(increments, 3, 3, threads, [](const float* in, int64_t steps, float* out) {
+// Re-express per-step rotation increments between world-frame (left) and
+// body-frame (right) multiplication, as scipy ``Rotation`` composes them.
+F32 rotvec_left_to_right_batch(const F64& increments, int threads) {
+  return map_rows(increments, 3, 3, threads, [](const double* in, int64_t steps, float* out) {
     Quat current{0.0, 0.0, 0.0, 1.0};
     for (int64_t t = 0; t < steps; ++t) {
       const double v[3] = {in[t * 3], in[t * 3 + 1], in[t * 3 + 2]};
@@ -762,8 +790,8 @@ F32 rotvec_left_to_right_batch(const F32& increments, int threads) {
   });
 }
 
-F32 rotvec_right_to_left_batch(const F32& increments, int threads) {
-  return map_rows(increments, 3, 3, threads, [](const float* in, int64_t steps, float* out) {
+F32 rotvec_right_to_left_batch(const F64& increments, int threads) {
+  return map_rows(increments, 3, 3, threads, [](const double* in, int64_t steps, float* out) {
     Quat current{0.0, 0.0, 0.0, 1.0};
     for (int64_t t = 0; t < steps; ++t) {
       const double v[3] = {in[t * 3], in[t * 3 + 1], in[t * 3 + 2]};
@@ -776,7 +804,7 @@ F32 rotvec_right_to_left_batch(const F32& increments, int threads) {
   });
 }
 
-}
+}  // namespace
 
 void register_resample(py::module_& module) {
   module.def("resample_uniform_batch", &resample_uniform_batch, py::arg("values"), py::arg("source_hz"),

actioncodec2_kernels/setbpe.cpp

@@ -1,3 +1,10 @@
+// Set-BPE serialization: quantized cell trajectories <-> token streams.
+//
+// A port of the runtime's SpatialVocabulary / TemporalVocabulary (Python,
+// ``tokenization/setbpe_{spatial,temporal}.py``).  Every merge is applied from
+// a priority queue ordered exactly as the Python heaps order their tuples, so
+// the token streams are identical, not merely equivalent.
+
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
@@ -5,6 +12,8 @@
 #include <algorithm>
 #include <cstdint>
 #include <functional>
+#include <memory>
+#include <mutex>
 #include <queue>
 #include <stdexcept>
 #include <string>
@@ -14,6 +23,8 @@
 #include <vector>
 
 #include "parallel.hpp"
+#include "repacking.hpp"
+#include "token_streams.hpp"
 
 namespace py = pybind11;
 
@@ -74,7 +85,7 @@ class SetBPE {
     for (int64_t token = 0; token < spatial_size_; ++token)
       spatial_anchor_[token] = lowest_bit(support_[token]);
 
-
+    // Rooted temporal tokens: a pattern of (time offset, spatial token) pairs.
     vocab_size_ = spatial_size_ + static_cast<int64_t>(temporal_merges.size());
     patterns_.resize(vocab_size_);
     root_.resize(vocab_size_);
@@ -113,9 +124,16 @@ class SetBPE {
   }
 
   int64_t vocab_size() const { return vocab_size_; }
-  int64_t dimension() const { return static_cast<int64_t>(bins_.size()); }
+  int64_t dimension() const { return bins_.size(); }
 
+  const ac2pack::Search& packing_search() const {
+    std::call_once(packing_once_, [&] {
+      packing_ = std::make_unique<ac2pack::Search>(bins_, support_, expansion_, patterns_, root_);
+    });
+    return *packing_;
+  }
 
+  // One dense [T, D] trajectory with -1 for semantically absent slots.
   std::vector<int64_t> encode(const int64_t* cells, int64_t steps) const {
     const int64_t dimension = static_cast<int64_t>(bins_.size());
     if (steps < 1) throw std::invalid_argument("every episode must have at least one step");
@@ -134,10 +152,10 @@ class SetBPE {
         throw std::invalid_argument("active dimensions must remain constant within an episode");
     }
 
-
+    // Instances, as in the Python encoder: token, anchor step, liveness.
     std::vector<int64_t> tokens, anchors;
     std::vector<char> active;
-    std::vector<std::vector<std::pair<int64_t, int64_t>>> at(steps);
+    std::vector<std::vector<std::pair<int64_t, int64_t>>> at(steps);  // (token, instance)
     for (int64_t t = 0; t < steps; ++t) {
       for (int64_t token : spatial_encode(cells + t * dimension)) {
         at[t].emplace_back(token, static_cast<int64_t>(tokens.size()));
@@ -147,7 +165,7 @@ class SetBPE {
       }
     }
 
-
+    // (rank, anchor, root anchor dimension, left instance, right instance, child)
     using Offer = std::tuple<int64_t, int64_t, int64_t, int64_t, int64_t, int64_t>;
     std::priority_queue<Offer, std::vector<Offer>, std::greater<Offer>> frontier;
     auto offer = [&](int64_t left_instance, int64_t right_instance) {
@@ -218,7 +236,7 @@ class SetBPE {
     return stream;
   }
 
-
+  // Tokens -> dense [T, D] cells (-1 outside ``full``), validating like the runtime.
   std::vector<int64_t> decode(const int64_t* stream, int64_t length, uint64_t full) const {
     const int64_t dimension = static_cast<int64_t>(bins_.size());
     std::vector<uint64_t> coverage{0};
@@ -264,7 +282,7 @@ class SetBPE {
   }
 
  private:
-
+  // Spatial Set-BPE of one dense row: merges in rank order, then anchor order.
   std::vector<int64_t> spatial_encode(const int64_t* row) const {
     const int64_t dimension = static_cast<int64_t>(bins_.size());
     std::vector<int64_t> present;
@@ -272,7 +290,7 @@ class SetBPE {
       if (row[d] >= 0) present.push_back(offsets_[d] + row[d]);
     }
     std::sort(present.begin(), present.end());
-    using Candidate = std::tuple<int64_t, int64_t, int64_t, int64_t>;
+    using Candidate = std::tuple<int64_t, int64_t, int64_t, int64_t>;  // rank, child, left, right
     std::priority_queue<Candidate, std::vector<Candidate>, std::greater<Candidate>> heap;
     for (size_t i = 0; i < present.size(); ++i) {
       for (size_t j = i + 1; j < present.size(); ++j) {
@@ -281,7 +299,7 @@ class SetBPE {
           heap.emplace(found->second.rank, found->second.child, present[i], present[j]);
       }
     }
-    std::vector<int64_t> tokens = present;
+    std::vector<int64_t> tokens = present;  // a small set; linear lookups are fastest
     auto contains = [&](int64_t token) {
       return std::find(tokens.begin(), tokens.end(), token) != tokens.end();
     };
@@ -325,13 +343,40 @@ class SetBPE {
   std::vector<std::vector<std::pair<int64_t, int64_t>>> patterns_;
   std::vector<int64_t> root_;
   std::unordered_map<std::tuple<int64_t, int64_t, int64_t>, Rule, TupleHash> temporal_rules_;
+  mutable std::once_flag packing_once_;
+  mutable std::unique_ptr<ac2pack::Search> packing_;
 };
 
+py::list encode_search_batch(const SetBPE& codec,
+    const py::array_t<int64_t, py::array::c_style | py::array::forcecast>& cells,
+    bool spatial, bool temporal, int window, int beam, int threads) {
+  if (cells.ndim() != 3 || cells.shape(0) < 1 || cells.shape(1) < 1 ||
+      cells.shape(2) != codec.dimension())
+    throw std::invalid_argument("expected nonempty [B,T,D] cells matching codec width");
+  if (window < 1 || window * codec.dimension() > 256 || beam < 1 || beam > 64)
+    throw std::invalid_argument("packing requires window*D <= 256 and beam in [1,64]");
+  const auto* search = spatial || temporal ? &codec.packing_search() : nullptr;
+  const int64_t batch = cells.shape(0), steps = cells.shape(1), width = cells.shape(2);
+  std::vector<std::vector<int64_t>> streams(batch);
+  {
+    py::gil_scoped_release release;
+    ac2::parallel_for(batch, threads, [&](int64_t index) {
+      const auto* row = cells.data() + index * steps * width;
+      auto original = codec.encode(row, steps);
+      streams[index] = search ? search->apply(row, steps, original, spatial, temporal, window, beam)
+                              : std::move(original);
+    });
+  }
+  py::list output(batch);
+  for (int64_t index = 0; index < batch; ++index) output[index] = py::cast(streams[index]);
+  return output;
+}
+
+// Encode B dense [T, D] trajectories in parallel; returns one stream per row.
 py::list encode_batch(const SetBPE& codec,
                       const py::array_t<int64_t, py::array::c_style | py::array::forcecast>& cells,
                       int threads) {
-  if (cells.ndim() != 3 || cells.shape(1) < 1 || cells.shape(2) != codec.dimension())
-    throw std::invalid_argument("cells must have shape [B, T, D] with the codec's dimension");
+  if (cells.ndim() != 3) throw std::invalid_argument("cells must have shape [B, T, D]");
   const int64_t batch = cells.shape(0), steps = cells.shape(1), width = cells.shape(2);
   const int64_t* data = cells.data();
   std::vector<std::vector<int64_t>> streams(batch);
@@ -346,15 +391,11 @@ py::list encode_batch(const SetBPE& codec,
   return output;
 }
 
+// Decode B equal-horizon streams to dense [B, T, D] cells (-1 outside ``full``).
 py::array_t<int64_t> decode_batch(const SetBPE& codec, const std::vector<std::vector<int64_t>>& streams,
                                   int64_t width, uint64_t full, int threads) {
   const int64_t batch = static_cast<int64_t>(streams.size());
   if (batch < 1) throw std::invalid_argument("streams must not be empty");
-  if (width != codec.dimension())
-    throw std::invalid_argument("width must equal the codec's dimension");
-  const uint64_t all = (uint64_t{1} << codec.dimension()) - 1;
-  if (!full || (full & ~all))
-    throw std::invalid_argument("full_support must select existing dimensions");
   std::vector<std::vector<int64_t>> cells(batch);
   {
     py::gil_scoped_release release;
@@ -375,7 +416,7 @@ py::array_t<int64_t> decode_batch(const SetBPE& codec, const std::vector<std::ve
   return output;
 }
 
-}
+}  // namespace
 
 void register_setbpe(py::module_& module) {
   py::class_<SetBPE>(module, "SetBPE")
@@ -386,38 +427,14 @@ void register_setbpe(py::module_& module) {
       .def_property_readonly("vocab_size", &SetBPE::vocab_size)
       .def("encode_batch", &encode_batch, py::arg("cells"), py::arg("threads") = 0,
            "Dense [B, T, D] cells (-1 = absent) -> one token stream per row.")
+      .def("encode_search_batch", &encode_search_batch, py::arg("cells"),
+           py::arg("spatial_search") = false, py::arg("temporal_search") = false,
+           py::arg("search_window") = 32, py::arg("search_beam") = 4, py::arg("threads") = 0)
       .def("decode_batch", &decode_batch, py::arg("streams"), py::arg("width"),
            py::arg("full_support"), py::arg("threads") = 0,
            "Equal-horizon streams -> dense [B, T, D] cells, -1 outside full_support.")
       .def("decode_checked_batch", [](const SetBPE& codec, const py::iterable& streams,
-                                     int64_t width, uint64_t full, int threads) {
-        const py::object numpy_integer = py::module_::import("numpy").attr("integer");
-        py::object integral;
-        std::vector<std::vector<int64_t>> rows;
-        if (PyList_Check(streams.ptr())) rows.reserve(PyList_GET_SIZE(streams.ptr()));
-        for (py::handle stream : streams) {
-          std::vector<int64_t> row;
-          if (PyList_Check(stream.ptr())) row.reserve(PyList_GET_SIZE(stream.ptr()));
-          for (py::handle raw : py::reinterpret_borrow<py::iterable>(stream)) {
-            bool integer = PyLong_Check(raw.ptr()) || py::isinstance(raw, numpy_integer);
-            if (!integer) {
-              if (!integral) integral = py::module_::import("numbers").attr("Integral");
-              integer = py::isinstance(raw, integral);
-            }
-            if (PyBool_Check(raw.ptr()) || !integer)
-              throw std::invalid_argument("tokens must be integers");
-            py::object number = PyLong_Check(raw.ptr())
-                ? py::reinterpret_borrow<py::object>(raw)
-                : py::reinterpret_steal<py::object>(PyNumber_Long(raw.ptr()));
-            if (!number) throw py::error_already_set();
-            int overflow = 0;
-            const int64_t token = PyLong_AsLongLongAndOverflow(number.ptr(), &overflow);
-            if (PyErr_Occurred()) throw py::error_already_set();
-            if (overflow) throw std::invalid_argument("token is outside the vocabulary");
-            row.push_back(token);
-          }
-          rows.push_back(std::move(row));
-        }
-        return decode_batch(codec, rows, width, full, threads);
+              int64_t width, uint64_t full, int threads) {
+        return decode_batch(codec, ac2::checked_token_streams(streams, codec.vocab_size()), width, full, threads);
       }, py::arg("streams"), py::arg("width"), py::arg("full_support"), py::arg("threads") = 0);
 }

actioncodec2_kernels/so3.cpp

@@ -1,3 +1,5 @@
+// Intrinsic first- and stateful second-order SO(3) error-feedback encoders.
+
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
 
@@ -104,6 +106,8 @@ struct SO3Layout {
   bool left;
   int32_t input_order;
   double projection_radius = std::numeric_limits<double>::infinity();
+  bool lazy = false;              // momentum_lazy ranking (first order only)
+  std::array<int32_t, 3> zero{};  // exact zero primitive per axis, used by ``lazy``
 };
 
 void project_primitive(double* value, double radius) {
@@ -150,22 +154,31 @@ std::vector<int32_t> smooth_axis(const SO3Layout& layout, int axis, double reque
   return output;
 }
 
-bool earlier(double error, const std::array<int32_t, 3>& choice, double best_error,
-             const std::array<int32_t, 3>& best_choice, bool found) {
-  return !found || error < best_error || (error == best_error && choice < best_choice);
-}
+// Lexicographic selection key: (moving axes, zero distance) under ``lazy``, then error, cell.
+struct ChoiceKey {
+  int moving = 0;
+  int distance = 0;
+  double error = 0.0;
+  std::array<int32_t, 3> choice{};
+  bool operator<(const ChoiceKey& o) const {
+    if (moving != o.moving) return moving < o.moving;
+    if (distance != o.distance) return distance < o.distance;
+    if (error != o.error) return error < o.error;
+    return choice < o.choice;
+  }
+};
 
 std::array<int32_t, 3> smooth_choice(const SO3Layout& layout,
                                      const std::array<double, 3>& request,
                                      const Quaternion& reconstruction,
                                      const Quaternion& target,
                                      const std::array<double, 3>& direction,
-                                     const std::array<int32_t, 3>& nearest) {
+                                     const std::array<int32_t, 3>& nearest,
+                                     const std::array<int32_t, 3>& previous = {-1, -1, -1}) {
   const auto x = smooth_axis(layout, 0, request[0], nearest[0]);
   const auto y = smooth_axis(layout, 1, request[1], nearest[1]);
   const auto z = smooth_axis(layout, 2, request[2], nearest[2]);
-  std::array<int32_t, 3> best_any{}, best_preferred{};
-  double any_error = 0.0, preferred_error = 0.0;
+  ChoiceKey best_any, best_preferred;
   bool found_any = false, found_preferred = false;
   for (const int32_t ix : x) {
     for (const int32_t iy : y) {
@@ -179,23 +192,30 @@ std::array<int32_t, 3> smooth_choice(const SO3Layout& layout,
         const double error = std::sqrt(remaining[0] * remaining[0] + remaining[1] * remaining[1] +
                                        remaining[2] * remaining[2]);
         if (error > layout.limit + 1e-12) continue;
-        if (earlier(error, choice, any_error, best_any, found_any)) {
-          best_any = choice;
-          any_error = error;
+        if (layout.lazy && choice == previous) return choice;
+        ChoiceKey key;
+        key.error = error;
+        key.choice = choice;
+        if (layout.lazy) {
+          for (int axis = 0; axis < 3; ++axis) {
+            key.moving += choice[axis] != layout.zero[axis];
+            key.distance += std::abs(choice[axis] - layout.zero[axis]);
+          }
+        }
+        if (!found_any || key < best_any) {
+          best_any = key;
           found_any = true;
         }
         const double dot = vector[0] * direction[0] + vector[1] * direction[1] +
                            vector[2] * direction[2];
-        if (dot >= -1e-15 && earlier(error, choice, preferred_error, best_preferred,
-                                  found_preferred)) {
-          best_preferred = choice;
-          preferred_error = error;
+        if (dot >= -1e-15 && (!found_preferred || key < best_preferred)) {
+          best_preferred = key;
           found_preferred = true;
         }
       }
     }
   }
-  return found_preferred ? best_preferred : (found_any ? best_any : nearest);
+  return found_preferred ? best_preferred.choice : (found_any ? best_any.choice : nearest);
 }
 
 void encode_trajectory(const SO3Layout& layout, const float* actions, int64_t steps,
@@ -203,14 +223,16 @@ void encode_trajectory(const SO3Layout& layout, const float* actions, int64_t st
   Quaternion truth{0.0, 0.0, 0.0, 1.0};
   Quaternion reconstruction = truth;
   std::array<double, 3> direction{};
+  std::array<int32_t, 3> previous = layout.zero;
   for (int64_t step = 0; step < steps; ++step) {
     const Quaternion value = from_input(actions + 3 * step, layout.rotation_vector);
     truth = layout.input_order == 0 ? value : apply(truth, value, layout.left);
     const auto request = as_rotvec(residual(truth, reconstruction, layout.left));
     auto choice = nearest_choice(layout, request);
     if (layout.smooth) {
-      choice = smooth_choice(layout, request, reconstruction, truth, direction, choice);
+      choice = smooth_choice(layout, request, reconstruction, truth, direction, choice, previous);
     }
+    previous = choice;
     const double vector[3] = {primitive(layout, 0, choice[0]),
                               primitive(layout, 1, choice[1]),
                               primitive(layout, 2, choice[2])};
@@ -222,7 +244,7 @@ void encode_trajectory(const SO3Layout& layout, const float* actions, int64_t st
   }
 }
 
-}
+}  // namespace
 
 py::array_t<int32_t> so3_second_order_encode_batch(
     const py::array_t<double, py::array::c_style | py::array::forcecast>& table,
@@ -263,7 +285,7 @@ py::array_t<int32_t> so3_second_order_encode_batch(
       const auto request = as_rotvec(correction);
       auto choice = nearest_choice(layout, request);
       if (smooth) {
-
+        // Python scores U^-1 W for both conventions; use exactly that order.
         SO3Layout selection = layout;
         selection.left = false;
         choice = smooth_choice(selection, request, identity, correction, direction, choice);
@@ -284,10 +306,11 @@ py::array_t<int32_t> so3_encode_batch(
     const py::array_t<double, py::array::c_style | py::array::forcecast>& table,
     const py::array_t<int32_t, py::array::c_style | py::array::forcecast>& bins,
     const py::array_t<float, py::array::c_style | py::array::forcecast>& actions,
-    int32_t input_order, bool rotation_vector, bool left, double limit, bool smooth,
+    int32_t input_order, bool rotation_vector, bool left, double limit, int32_t selection,
     int threads) {
   if (table.ndim() != 2 || table.shape(0) != 3 || bins.size() != 3 || actions.ndim() != 3 ||
-      actions.shape(2) != 3 || (input_order != 0 && input_order != 1) || limit <= 0.0) {
+      actions.shape(2) != 3 || (input_order != 0 && input_order != 1) || limit <= 0.0 ||
+      selection < 0 || selection > 2) {
     throw std::invalid_argument("invalid SO(3) batch encoder arguments");
   }
   for (int axis = 0; axis < 3; ++axis) {
@@ -295,11 +318,28 @@ py::array_t<int32_t> so3_encode_batch(
       throw std::invalid_argument("SO(3) primitive count exceeds table width");
     }
   }
-  const SO3Layout layout{table.data(), bins.data(), table.shape(1), limit, smooth,
-                         rotation_vector, left, input_order};
+  // selection: 0 nearest, 1 smooth, 2 momentum_lazy
+  SO3Layout layout{table.data(), bins.data(), table.shape(1), limit, selection != 0,
+                   rotation_vector, left, input_order};
+  layout.lazy = selection == 2;
+  for (int axis = 0; axis < 3; ++axis) {
+    double best = std::numeric_limits<double>::infinity();
+    for (int32_t index = 0; index < bins.data()[axis]; ++index) {
+      if (std::fabs(primitive(layout, axis, index)) < best) {
+        best = std::fabs(primitive(layout, axis, index));
+        layout.zero[axis] = index;
+      }
+    }
+    if (layout.lazy && best != 0.0) {
+      throw std::invalid_argument("momentum_lazy needs an exact zero primitive per axis");
+    }
+  }
   const int64_t batch = actions.shape(0);
   const int64_t steps = actions.shape(1);
   py::array_t<int32_t> output({batch, steps, int64_t{3}});
+#ifndef _OPENMP
+  (void)threads;
+#endif
   {
     py::gil_scoped_release release;
     ac2::parallel_for(batch, threads, [&](int64_t index) {

actioncodec2_kernels/sparse.cpp

@@ -0,0 +1,1308 @@
+
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <algorithm>
+#include <array>
+#include <chrono>
+#include <cstdint>
+#include <limits>
+#include <queue>
+#include <set>
+#include <stdexcept>
+#include <string>
+#include <tuple>
+#include <unordered_map>
+#include <vector>
+
+#include "parallel.hpp"
+
+#ifdef _OPENMP
+#include <omp.h>
+#endif
+#if defined(_OPENMP) && defined(__GLIBCXX__)
+#include <parallel/algorithm>
+#define AC2_PARALLEL_SORT __gnu_parallel::sort
+#else
+#define AC2_PARALLEL_SORT std::sort
+#endif
+
+namespace py = pybind11;
+
+namespace {
+
+constexpr int kMaxLookahead = 32;
+constexpr size_t kCertifyBatch = 32;
+constexpr size_t kApplyGrain = 256;  // minimum candidate pairs per parallel apply group
+
+constexpr size_t kShards = 64;  // rule-table shards updated concurrently after a merge
+
+inline uint64_t mix(uint64_t key) { return key * 0x9E3779B97F4A7C15ull; }
+inline size_t shard_of(uint64_t key) { return static_cast<size_t>(mix(key) >> 58); }
+
+// Small open-addressing accumulator for rule-count deltas; clearing costs O(used).
+class DeltaMap {
+ public:
+  void add(uint64_t key, int64_t delta) {
+    if ((filled_.size() + 1) * 2 > keys_.size()) grow();
+    size_t h = slot(key);
+    while (keys_[h] != kEmpty && keys_[h] != key) h = (h + 1) & (keys_.size() - 1);
+    if (keys_[h] == kEmpty) {
+      keys_[h] = key;
+      filled_.push_back(static_cast<uint32_t>(h));
+    }
+    values_[h] += delta;
+  }
+  template <class F>
+  void each(F&& visit) const {
+    for (uint32_t h : filled_) visit(keys_[h], values_[h]);
+  }
+  void clear() {
+    for (uint32_t h : filled_) {
+      keys_[h] = kEmpty;
+      values_[h] = 0;
+    }
+    filled_.clear();
+  }
+
+ private:
+  static constexpr uint64_t kEmpty = ~0ull;
+  size_t slot(uint64_t key) const {
+    return static_cast<size_t>(mix(key) >> 20) & (keys_.size() - 1);
+  }
+  void grow() {
+    std::vector<uint64_t> keys = std::move(keys_);
+    std::vector<int64_t> values = std::move(values_);
+    std::vector<uint32_t> filled = std::move(filled_);
+    const size_t capacity = std::max<size_t>(64, keys.size() * 2);
+    keys_.assign(capacity, kEmpty);
+    values_.assign(capacity, 0);
+    filled_.clear();
+    for (uint32_t h : filled) add(keys[h], values[h]);
+  }
+  std::vector<uint64_t> keys_;
+  std::vector<int64_t> values_;
+  std::vector<uint32_t> filled_;
+};
+
+inline uint64_t rule_key(int32_t left, int32_t right, int32_t offset) {
+  return (static_cast<uint64_t>(left) << 40) | (static_cast<uint64_t>(right) << 16) |
+         static_cast<uint64_t>(offset);
+}
+
+inline void unpack(uint64_t key, int32_t& left, int32_t& right, int32_t& offset) {
+  left = static_cast<int32_t>(key >> 40);
+  right = static_cast<int32_t>((key >> 16) & 0xffffffull);
+  offset = static_cast<int32_t>(key & 0xffffull);
+}
+
+// Immutable-per-token geometry shared by the trainer, encoder and decoder.
+struct Tables {
+  int32_t dims = 0;       // physical slots including E
+  int32_t e_dim = 0;      // index of E
+  int32_t lookahead = 0;  // maximum token span in frames
+  std::vector<uint32_t> masks;                 // token * lookahead -> support mask
+  std::vector<uint32_t> emask;                 // token -> bit per offset holding E
+  std::vector<int8_t> first_from;              // token * lookahead + k -> first support offset >= k
+  std::vector<int32_t> span, rdim;
+  std::vector<std::vector<std::pair<int32_t, int32_t>>> pattern;  // (offset, spatial token)
+  int32_t base = 0;
+
+  void init(const std::vector<uint32_t>& spatial_support, int32_t dims_, int32_t lookahead_) {
+    dims = dims_;
+    e_dim = dims_ - 1;
+    lookahead = lookahead_;
+    if (lookahead < 1 || lookahead > kMaxLookahead) throw std::invalid_argument("lookahead");
+    if (dims < 2 || dims > 16) throw std::invalid_argument("dims must be in [2, 16]");
+    base = static_cast<int32_t>(spatial_support.size());
+    for (int32_t t = 0; t < base; ++t) {
+      uint32_t s = spatial_support[t];
+      if (!s) throw std::invalid_argument("spatial token without support");
+      push({{0, t}}, std::vector<uint32_t>{s});
+    }
+  }
+
+  void push(std::vector<std::pair<int32_t, int32_t>> pat, const std::vector<uint32_t>& m) {
+    const size_t token = span.size();
+    masks.resize((token + 1) * lookahead, 0);
+    uint32_t e = 0;
+    int32_t sp = 0;
+    for (size_t o = 0; o < m.size(); ++o) {
+      masks[token * lookahead + o] = m[o];
+      if (m[o]) sp = static_cast<int32_t>(o) + 1;
+      if (m[o] >> e_dim & 1u) e |= 1u << o;
+    }
+    emask.push_back(e);
+    first_from.resize((token + 1) * lookahead, -1);
+    for (int32_t k = sp - 1, next = -1; k >= 0; --k) {
+      if (masks[token * lookahead + k]) next = k;
+      first_from[token * lookahead + k] = static_cast<int8_t>(next);
+    }
+    span.push_back(sp);
+    rdim.push_back(__builtin_ctz(m[0]));
+    pattern.push_back(std::move(pat));
+  }
+
+  int32_t size() const { return static_cast<int32_t>(span.size()); }
+  uint32_t mask(int32_t token, int32_t o) const {
+    return masks[static_cast<size_t>(token) * lookahead + o];
+  }
+
+  // Merge geometry; returns false when the union would exceed the lookahead or overlap.
+  bool compose(int32_t a, int32_t b, int32_t off, std::vector<uint32_t>& out) const {
+    if (off < 0 || off >= lookahead) return false;
+    const int32_t sp = std::max(span[a], off + span[b]);
+    if (sp > lookahead) return false;
+    out.assign(sp, 0);
+    for (int32_t o = 0; o < span[a]; ++o) out[o] = mask(a, o);
+    for (int32_t o = 0; o < span[b]; ++o) {
+      if (out[o + off] & mask(b, o)) return false;
+      out[o + off] |= mask(b, o);
+    }
+    return true;
+  }
+
+  int32_t add_merge(int32_t a, int32_t b, int32_t off) {
+    std::vector<uint32_t> m;
+    if (!compose(a, b, off, m)) throw std::invalid_argument("illegal temporal merge");
+    auto pat = pattern[a];
+    for (auto [o, t] : pattern[b]) pat.push_back({o + off, t});
+    std::sort(pat.begin(), pat.end());
+    push(std::move(pat), m);
+    return size() - 1;
+  }
+};
+
+using Key = std::tuple<int64_t, int32_t, int64_t, int32_t>;
+using InstanceId = uint32_t;
+constexpr InstanceId kNoInstance = UINT32_MAX;
+
+// Mutable instance state over a global frame axis.  Used per chunk by the encoder and
+// globally by the trainer.
+struct Instances {
+  const Tables* tb = nullptr;
+  int64_t frames = 0;
+  std::vector<int32_t> tok;
+  std::vector<int64_t> root;
+  std::vector<uint8_t> alive;
+  std::vector<InstanceId> member;  // frame * dims + dim -> instance
+
+  void reset(const Tables* tables, int64_t frame_count) {
+    tb = tables;
+    frames = frame_count;
+    tok.clear();
+    root.clear();
+    alive.clear();
+    member.assign(static_cast<size_t>(frame_count) * tables->dims, kNoInstance);
+  }
+
+  InstanceId cell(int64_t f, int32_t d) const {
+    return member[static_cast<size_t>(f) * tb->dims + d];
+  }
+
+  Key key(InstanceId i) const {
+    const uint32_t e = tb->emask[tok[i]];
+    return Key(root[i], e ? 1 : 0, e ? root[i] + __builtin_ctz(e) : 0, tb->rdim[tok[i]]);
+  }
+
+  InstanceId add(int32_t token, int64_t r) {
+    if (tok.size() >= kNoInstance) throw std::length_error("instance ids");
+    const InstanceId i = static_cast<InstanceId>(tok.size());
+    tok.push_back(token);
+    root.push_back(r);
+    alive.push_back(1);
+    place(i);
+    return i;
+  }
+
+  // Reserves dead slots that disjoint frame ranges then fill concurrently via ``place``.
+  InstanceId reserve(size_t count) {
+    const size_t base = tok.size();
+    if (base + count > static_cast<size_t>(kNoInstance)) throw std::length_error("instance ids");
+    tok.resize(base + count, 0);
+    root.resize(base + count, 0);
+    alive.resize(base + count, 0);
+    return static_cast<InstanceId>(base);
+  }
+
+  void place(InstanceId i) {
+    const int32_t token = tok[i];
+    const int64_t r = root[i];
+    alive[i] = 1;
+    for (int32_t o = 0; o < tb->span[token]; ++o) {
+      uint32_t m = tb->mask(token, o);
+      while (m) {
+        const int32_t d = __builtin_ctz(m);
+        m &= m - 1;
+        InstanceId& slot = member[static_cast<size_t>(r + o) * tb->dims + d];
+        if (slot != kNoInstance) throw std::runtime_error("instance overlap");
+        slot = i;
+      }
+    }
+  }
+
+  void kill(InstanceId i) {
+    alive[i] = 0;
+    const int32_t token = tok[i];
+    for (int32_t o = 0; o < tb->span[token]; ++o) {
+      uint32_t m = tb->mask(token, o);
+      while (m) {
+        const int32_t d = __builtin_ctz(m);
+        m &= m - 1;
+        member[static_cast<size_t>(root[i] + o) * tb->dims + d] = kNoInstance;
+      }
+    }
+  }
+
+  bool first(InstanceId a, InstanceId b) const {
+    return std::make_pair(root[a], tb->rdim[tok[a]]) < std::make_pair(root[b], tb->rdim[tok[b]]);
+  }
+
+  // Ordered candidate pairs (a, b, offset) involving i inside the lookahead window.
+  // Each neighbour is emitted once, at its first support cell inside the window.
+  template <class F>
+  void partners(InstanceId i, bool allow_zero, F&& emit) const {
+    const int64_t r = root[i];
+    const int32_t L = tb->lookahead;
+    const int64_t lo = std::max<int64_t>(0, r - L + 1), hi = std::min<int64_t>(frames - 1, r + L - 1);
+    for (int64_t f = lo; f <= hi; ++f) {
+      for (int32_t d = 0; d < tb->dims; ++d) {
+        const InstanceId j = cell(f, d);
+        if (j == kNoInstance || j == i) continue;
+        const int32_t tj = tok[j];
+        const int64_t o = f - root[j];
+        const int64_t k = std::max<int64_t>(0, lo - root[j]);
+        if (tb->first_from[static_cast<size_t>(tj) * L + k] != o) continue;
+        if (__builtin_ctz(tb->mask(tj, static_cast<int32_t>(o))) != d) continue;
+        const bool ij = first(i, j);
+        const InstanceId a = ij ? i : j, b = ij ? j : i;
+        const int64_t off = root[b] - root[a];
+        if (off >= L || (off == 0 && !allow_zero)) continue;
+        if (std::max<int64_t>(tb->span[tok[a]], off + tb->span[tok[b]]) > L) continue;
+        emit(a, b, static_cast<int32_t>(off));
+      }
+    }
+  }
+
+  // Would merging a (rooted first) with b keep the stream decodable?
+  bool valid(InstanceId a, InstanceId b, int32_t off) const {
+    const int32_t ta = tok[a], tb_ = tok[b];
+    const int32_t sp = std::max(tb->span[ta], off + tb->span[tb_]);
+    uint32_t emask = tb->emask[ta] | (tb->emask[tb_] << off);
+    const Key km(root[a], emask ? 1 : 0, emask ? root[a] + __builtin_ctz(emask) : 0,
+                 tb->rdim[ta]);
+    for (int32_t o = 0; o < sp; ++o) {
+      uint32_t m = (o < tb->span[ta] ? tb->mask(ta, o) : 0) |
+                   (o >= off && o - off < tb->span[tb_] ? tb->mask(tb_, o - off) : 0);
+      if (!m) continue;
+      const int64_t u = root[a] + o;
+      const InstanceId h = cell(u, tb->e_dim);
+      if (h != kNoInstance && h != a && h != b && !(key(h) > km)) return false;
+      if (emask >> o & 1u) {
+        for (int32_t d = 0; d < tb->dims; ++d) {
+          const InstanceId j = cell(u, d);
+          if (j != kNoInstance && j != a && j != b && !(key(j) < km)) return false;
+        }
+      }
+    }
+    return true;
+  }
+
+  // Right partner of left instance a under rule (left, right, off), or -1.
+  InstanceId right_of(InstanceId a, int32_t right, int32_t off) const {
+    const int64_t f = root[a] + off;
+    if (f >= frames) return kNoInstance;
+    const InstanceId b = cell(f, tb->rdim[right]);
+    if (b == kNoInstance || b == a || !alive[b] || tok[b] != right || root[b] != f) return kNoInstance;
+    return b;
+  }
+
+  // Left partner of right instance b under rule (left, right, off), or -1.
+  InstanceId left_of(InstanceId b, int32_t left, int32_t off) const {
+    const int64_t f = root[b] - off;
+    if (f < 0) return kNoInstance;
+    const InstanceId a = cell(f, tb->rdim[left]);
+    if (a == kNoInstance || a == b || !alive[a] || tok[a] != left || root[a] != f) return kNoInstance;
+    return a;
+  }
+
+  std::vector<int32_t> serialize() const {
+    std::vector<InstanceId> live;
+    for (size_t i = 0; i < tok.size(); ++i)
+      if (alive[i]) live.push_back(static_cast<InstanceId>(i));
+    std::sort(live.begin(), live.end(), [&](InstanceId x, InstanceId y) { return key(x) < key(y); });
+    std::vector<int32_t> out;
+    out.reserve(live.size());
+    for (InstanceId i : live) out.push_back(tok[i]);
+    return out;
+  }
+};
+
+struct Stat {
+  bool retired = false;  // a rule is learned at most once; leftovers stay unmerged
+  int64_t raw = 0;
+  uint32_t version = 0;
+  int64_t exact = -1;  // matching- and validity-certified count at exact_version
+  uint32_t exact_version = 0;
+  int64_t added = 0;  // cumulative pair insertions; bounds growth since certification
+  int64_t added_at_exact = 0;
+  uint32_t dirty_epoch = 0;  // merge batch that last queued this key for re-push
+  bool live = false;         // a heap entry (queued_version, queued_count) represents the key
+  uint32_t queued_version = 0;
+  int64_t queued_count = 0;
+};
+
+// Open-addressing rule table; references are invalidated only by inserting a new key.
+class StatMap {
+ public:
+  Stat* find(uint64_t key) {
+    if (keys_.empty()) return nullptr;
+    for (size_t h = slot(key);; h = (h + 1) & (keys_.size() - 1)) {
+      if (keys_[h] == key) return &values_[h];
+      if (keys_[h] == kEmpty) return nullptr;
+    }
+  }
+  Stat& operator[](uint64_t key) {
+    if ((used_ + 1) * 2 > keys_.size()) grow();
+    size_t h = slot(key);
+    while (keys_[h] != kEmpty && keys_[h] != key) h = (h + 1) & (keys_.size() - 1);
+    if (keys_[h] == kEmpty) {
+      keys_[h] = key;
+      ++used_;
+    }
+    return values_[h];
+  }
+  size_t size() const { return used_; }
+  template <class F>
+  void each(F&& visit) const {
+    for (size_t h = 0; h < keys_.size(); ++h)
+      if (keys_[h] != kEmpty) visit(keys_[h], values_[h]);
+  }
+
+ private:
+  static constexpr uint64_t kEmpty = ~0ull;
+  size_t slot(uint64_t key) const {
+    return static_cast<size_t>(mix(key) >> 20) & (keys_.size() - 1);
+  }
+  void grow() {
+    std::vector<uint64_t> keys = std::move(keys_);
+    std::vector<Stat> values = std::move(values_);
+    const size_t capacity = std::max<size_t>(1 << 16, keys.size() * 2);
+    keys_.assign(capacity, kEmpty);
+    values_.assign(capacity, Stat{});
+    used_ = 0;
+    for (size_t h = 0; h < keys.size(); ++h)
+      if (keys[h] != kEmpty) (*this)[keys[h]] = values[h];
+  }
+  std::vector<uint64_t> keys_;
+  std::vector<Stat> values_;
+  size_t used_ = 0;
+};
+
+// Rule table split by key hash so disjoint shards can be updated concurrently.
+class ShardedStats {
+ public:
+  Stat* find(uint64_t key) { return shards_[shard_of(key)].find(key); }
+  Stat& operator[](uint64_t key) { return shards_[shard_of(key)][key]; }
+  StatMap& shard(size_t p) { return shards_[p]; }
+  size_t size() const {
+    size_t n = 0;
+    for (const auto& s : shards_) n += s.size();
+    return n;
+  }
+  template <class F>
+  void each(F&& visit) const {
+    for (const auto& s : shards_) s.each(visit);
+  }
+
+ private:
+  StatMap shards_[kShards];
+};
+
+struct Entry {
+  int64_t count;
+  uint64_t key;
+  uint32_t version;
+  bool operator<(const Entry& o) const {
+    if (count != o.count) return count < o.count;
+    return key > o.key;
+  }
+};
+
+class Trainer {
+ public:
+  Trainer(std::vector<uint32_t> spatial_support, int32_t dims, int32_t lookahead,
+          int64_t vocab_size, bool allow_zero, int64_t min_count, int threads,
+          bool dense_census = false, int64_t max_instances = UINT32_MAX)
+      : vocab_size_(vocab_size), allow_zero_(allow_zero), dense_census_(dense_census),
+        min_count_(min_count), max_instances_(max_instances) {
+    if (max_instances_ < 1 || max_instances_ > UINT32_MAX)
+      throw std::invalid_argument("max_instances");
+    tables_.init(spatial_support, dims, lookahead);
+#ifdef _OPENMP
+    threads_ = threads > 0 ? threads : std::min(omp_get_max_threads(), 32);
+#else
+    (void)threads;
+    threads_ = 1;
+#endif
+  }
+
+  // anchors: global frame index per instance (chunks already separated by >= lookahead).
+  void set_instances(py::array_t<int64_t, py::array::c_style> anchors,
+                     py::array_t<int32_t, py::array::c_style> tokens, int64_t frames) {
+    auto a = anchors.unchecked<1>();
+    auto t = tokens.unchecked<1>();
+    if (a.shape(0) != t.shape(0)) throw std::invalid_argument("anchors/tokens length");
+    if (a.shape(0) > max_instances_)
+      throw std::length_error("live instance capacity exceeded before training");
+    st_.reset(&tables_, frames);
+    // Pre-sized so references into the outer vector stay valid while children are added.
+    occurrence_.assign(static_cast<size_t>(vocab_size_) + 16, {});
+    dead_.assign(static_cast<size_t>(vocab_size_) + 16, 0);
+    for (py::ssize_t i = 0; i < a.shape(0); ++i) {
+      const int32_t tk = t(i);
+      if (tk < 0 || tk >= tables_.base) throw std::invalid_argument("token outside spatial vocab");
+      occurrence_[tk].push_back(st_.add(tk, a(i)));
+    }
+    for (auto& occ : occurrence_)
+      std::sort(occ.begin(), occ.end(), [&](InstanceId x, InstanceId y) { return st_.root[x] < st_.root[y]; });
+    initial_instances_ = a.shape(0);
+    // Supports never change, so frames preceded by >= lookahead empty frames split the
+    // axis into ranges whose merges never read or write each other's cells.
+    safe_.clear();
+    int64_t empty = tables_.lookahead;
+    for (int64_t f = 0; f < frames; ++f) {
+      bool occupied = false;
+      for (int32_t d = 0; d < tables_.dims && !occupied; ++d) occupied = st_.cell(f, d) != kNoInstance;
+      if (!occupied) {
+        ++empty;
+        continue;
+      }
+      if (empty >= tables_.lookahead) safe_.push_back(f);
+      empty = 0;
+    }
+  }
+
+  py::dict train(bool progress) {
+    const auto t0 = std::chrono::steady_clock::now();
+    {
+      py::gil_scoped_release release;
+      census();
+      std::vector<uint64_t> keys;
+      keys.reserve(stats_.size());
+      stats_.each([&](uint64_t k, const Stat&) { keys.push_back(k); });
+      for (uint64_t k : keys) queue(k, stats_[k], stats_[k].raw);
+      initial_types_ = stats_.size();
+      census_seconds_ = seconds_since(t0);
+      run(progress);
+    }
+    py::dict report;
+    report["initial_instances"] = initial_instances_;
+    report["final_instances"] = live_;
+    report["initial_rule_types"] = initial_types_;
+    report["merges"] = static_cast<int64_t>(merges_.size());
+    report["dry_runs"] = dry_runs_;
+    report["rejected_by_validity"] = rejected_;
+    report["applied"] = applied_;
+    report["census_seconds"] = census_seconds_;
+    report["dry_run_seconds"] = dry_seconds_;
+    report["apply_seconds"] = apply_seconds_;
+    report["collect_seconds"] = collect_seconds_;
+    report["parallel_seconds"] = parallel_seconds_;
+    report["serial_seconds"] = serial_seconds_;
+    report["merge_seconds"] = merge_seconds_;
+    report["changes"] = changes_;
+    report["compactions"] = compactions_;
+    report["total_seconds"] = seconds_since(t0);
+    return report;
+  }
+
+  std::vector<std::tuple<int32_t, int32_t, int32_t>> merges() const { return merges_; }
+
+  // Count every forward pair once: blockwise parallel key generation, parallel sort, RLE.
+  void census() {
+    const int64_t n = static_cast<int64_t>(st_.tok.size());
+    const int64_t base = tables_.base;
+    const int64_t universe = base * base * tables_.lookahead;
+    if (dense_census_ && universe <= 4000000) {
+      std::vector<std::vector<int64_t>> local(
+          threads_, std::vector<int64_t>(static_cast<size_t>(universe), 0));
+#ifdef _OPENMP
+#pragma omp parallel num_threads(threads_)
+#endif
+      {
+#ifdef _OPENMP
+        auto& counts = local[omp_get_thread_num()];
+#else
+        auto& counts = local[0];
+#endif
+#ifdef _OPENMP
+#pragma omp for schedule(static)
+#endif
+        for (int64_t i = 0; i < n; ++i) {
+          const InstanceId ii = static_cast<InstanceId>(i);
+          st_.partners(ii, allow_zero_, [&](InstanceId a, InstanceId b, int32_t off) {
+            if (a == ii) {
+              const int64_t key = (static_cast<int64_t>(st_.tok[a]) * base + st_.tok[b]) *
+                                      tables_.lookahead + off;
+              ++counts[static_cast<size_t>(key)];
+            }
+          });
+        }
+      }
+      for (int64_t key = 0; key < universe; ++key) {
+        int64_t count = 0;
+        for (const auto& counts : local) count += counts[static_cast<size_t>(key)];
+        if (!count) continue;
+        const int32_t offset = static_cast<int32_t>(key % tables_.lookahead);
+        const int64_t pair = key / tables_.lookahead;
+        const int32_t left = static_cast<int32_t>(pair / base);
+        const int32_t right = static_cast<int32_t>(pair % base);
+        stats_[rule_key(left, right, offset)].raw += count;
+      }
+      return;
+    }
+    const int64_t block = int64_t{1} << 22;
+    std::vector<std::vector<uint64_t>> local(threads_);
+    std::vector<uint64_t> keys;
+    for (int64_t start = 0; start < n; start += block) {
+      const int64_t stop = std::min(n, start + block);
+#ifdef _OPENMP
+#pragma omp parallel num_threads(threads_)
+#endif
+      {
+#ifdef _OPENMP
+        auto& out = local[omp_get_thread_num()];
+#else
+        auto& out = local[0];
+#endif
+        out.clear();
+#ifdef _OPENMP
+#pragma omp for schedule(static)
+#endif
+        for (int64_t i = start; i < stop; ++i) {
+          const InstanceId ii = static_cast<InstanceId>(i);
+          st_.partners(ii, allow_zero_, [&](InstanceId a, InstanceId b, int32_t off) {
+            if (a == ii) out.push_back(rule_key(st_.tok[a], st_.tok[b], off));
+          });
+        }
+      }
+      keys.clear();
+      for (auto& v : local) keys.insert(keys.end(), v.begin(), v.end());
+      AC2_PARALLEL_SORT(keys.begin(), keys.end());
+      for (size_t k = 0; k < keys.size();) {
+        size_t e = k;
+        while (e < keys.size() && keys[e] == keys[k]) ++e;
+        stats_[keys[k]].raw += static_cast<int64_t>(e - k);
+        k = e;
+      }
+    }
+  }
+
+ private:
+  // Reclaim dead slots without changing token order, coverage, or candidate statistics.
+  void compact_instances() {
+    std::vector<InstanceId> remap(st_.tok.size(), kNoInstance);
+    std::vector<int32_t> tokens;
+    std::vector<int64_t> roots;
+    std::vector<uint8_t> alive;
+    tokens.reserve(static_cast<size_t>(live_));
+    roots.reserve(static_cast<size_t>(live_));
+    alive.reserve(static_cast<size_t>(live_));
+    for (size_t old = 0; old < st_.tok.size(); ++old) {
+      if (!st_.alive[old]) continue;
+      remap[old] = static_cast<InstanceId>(tokens.size());
+      tokens.push_back(st_.tok[old]);
+      roots.push_back(st_.root[old]);
+      alive.push_back(1);
+    }
+    if (tokens.size() != static_cast<size_t>(live_))
+      throw std::runtime_error("instance compaction live count mismatch");
+    for (auto& occ : occurrence_) {
+      size_t next = 0;
+      for (InstanceId old : occ) {
+        const InstanceId mapped = remap[static_cast<size_t>(old)];
+        if (mapped != kNoInstance) occ[next++] = mapped;
+      }
+      occ.resize(next);
+    }
+    for (auto& slot : st_.member)
+      if (slot != kNoInstance) slot = remap[static_cast<size_t>(slot)];
+    st_.tok.swap(tokens);
+    st_.root.swap(roots);
+    st_.alive.swap(alive);
+    std::fill(dead_.begin(), dead_.end(), 0);
+    ++compactions_;
+  }
+
+  void change(uint64_t k, int64_t delta, StatMap& table, std::vector<uint64_t>& dirty) {
+    Stat& s = table[k];
+    s.raw += delta;
+    if (delta > 0) s.added += delta;
+    ++s.version;
+    if (s.dirty_epoch != epoch_) {
+      s.dirty_epoch = epoch_;
+      dirty.push_back(k);
+    }
+  }
+
+  std::vector<InstanceId>& lefts(int32_t left) {
+    auto& occ = occurrence_[left];
+    // Lists are root-sorted by construction: atoms are sorted once, and every child list
+    // is appended in the root order of its single merge batch.
+    if (dead_[left] * 4 > static_cast<int64_t>(occ.size())) {
+      occ.erase(std::remove_if(occ.begin(), occ.end(), [&](InstanceId i) { return !st_.alive[i]; }),
+                occ.end());
+      dead_[left] = 0;
+    }
+    return occ;
+  }
+
+  static double seconds_since(std::chrono::steady_clock::time_point t) {
+    return std::chrono::duration<double>(std::chrono::steady_clock::now() - t).count();
+  }
+
+  // Candidate (a, b) pairs are visited in increasing root order of a by scanning the
+  // shorter occurrence list; the partner is one cell lookup away in either direction.
+  const std::vector<InstanceId>& scanned(int32_t left, int32_t right, bool& forward) const {
+    const auto& lo = occurrence_[left];
+    const auto& ro = occurrence_[right];
+    forward = left == right || lo.size() <= ro.size();
+    return forward ? lo : ro;
+  }
+
+  template <class F>
+  void pairs_in(int32_t left, int32_t right, int32_t off, size_t begin, size_t end,
+                F&& visit) const {
+    bool forward;
+    const auto& scan = scanned(left, right, forward);
+    for (size_t i = begin; i < end; ++i) {
+      const InstanceId x = scan[i];
+      if (!st_.alive[x]) continue;
+      if (forward) {
+        const InstanceId b = st_.right_of(x, right, off);
+        if (b != kNoInstance) visit(x, b);
+      } else {
+        const InstanceId a = st_.left_of(x, left, off);
+        if (a != kNoInstance) visit(a, x);
+      }
+    }
+  }
+
+  // Cuts of a root-sorted list into blocks of about ``grain`` entries at safe frames, so
+  // per-block replay (self-rule chains included) equals one sequential pass.
+  std::vector<size_t> blocks(const std::vector<InstanceId>& scan, size_t grain) const {
+    std::vector<size_t> cuts{0};
+    for (size_t next = grain; next < scan.size();) {
+      auto frame = std::upper_bound(safe_.begin(), safe_.end(), st_.root[scan[next]]);
+      if (frame == safe_.end()) break;
+      const size_t at = static_cast<size_t>(
+          std::lower_bound(scan.begin() + next, scan.end(), *frame,
+                           [&](InstanceId i, int64_t f) { return st_.root[i] < f; }) -
+          scan.begin());
+      if (at >= scan.size()) break;
+      cuts.push_back(at);
+      next = at + grain;
+    }
+    cuts.push_back(scan.size());
+    return cuts;
+  }
+
+  // Candidate pairs in scan order, gathered by blocks in parallel.
+  std::vector<std::pair<InstanceId, InstanceId>> collect(int32_t left, int32_t right, int32_t off) const {
+    bool forward;
+    const std::vector<size_t> cuts = blocks(scanned(left, right, forward), 1 << 16);
+    std::vector<std::vector<std::pair<InstanceId, InstanceId>>> parts(cuts.size() - 1);
+#ifdef _OPENMP
+#pragma omp parallel for schedule(dynamic, 1) num_threads(threads_) if (parts.size() > 1)
+#endif
+    for (size_t k = 0; k < parts.size(); ++k)
+      pairs_in(left, right, off, cuts[k], cuts[k + 1],
+               [&](InstanceId a, InstanceId b) { parts[k].emplace_back(a, b); });
+    std::vector<std::pair<InstanceId, InstanceId>> pairs;
+    for (auto& part : parts) pairs.insert(pairs.end(), part.begin(), part.end());
+    return pairs;
+  }
+
+  // Exact count of one block under the real merge's matching and validity tests.
+  int64_t count_exact(int32_t left, int32_t right, int32_t off, size_t begin, size_t end) const {
+    int64_t n = 0;
+    // Self rules chain: a may be the right half of a recent match (roots increase).
+    std::vector<InstanceId> recent;
+    size_t head = 0;
+    pairs_in(left, right, off, begin, end, [&](InstanceId a, InstanceId b) {
+      if (left == right) {
+        while (head < recent.size() && st_.root[recent[head]] < st_.root[a]) ++head;
+        if (std::find(recent.begin() + head, recent.end(), a) != recent.end()) return;
+      }
+      if (!st_.valid(a, b, off)) return;
+      if (left == right) recent.push_back(b);
+      ++n;
+    });
+    return n;
+  }
+
+  // Certifies a batch of keys; every key is split into blocks so one large list does
+  // not serialize the batch.  Read-only on shared state.
+  void certify(std::vector<Entry>& batch) {
+    struct Timer {
+      double& sink;
+      std::chrono::steady_clock::time_point t;
+      ~Timer() { sink += seconds_since(t); }
+    } timer{dry_seconds_, std::chrono::steady_clock::now()};
+    dry_runs_ += static_cast<int64_t>(batch.size());
+    struct Task {
+      size_t key, begin, end;
+    };
+    std::vector<Task> tasks;
+    for (size_t k = 0; k < batch.size(); ++k) {
+      int32_t l, r, o;
+      unpack(batch[k].key, l, r, o);
+      lefts(l);  // compaction mutates lists: do it before the parallel region
+      lefts(r);
+      bool forward;
+      const std::vector<size_t> cuts = blocks(scanned(l, r, forward), 1 << 15);
+      for (size_t c = 0; c + 1 < cuts.size(); ++c) tasks.push_back({k, cuts[c], cuts[c + 1]});
+    }
+    std::vector<int64_t> partial(tasks.size());
+#ifdef _OPENMP
+#pragma omp parallel for schedule(dynamic, 1) num_threads(threads_)
+#endif
+    for (size_t t = 0; t < tasks.size(); ++t) {
+      int32_t l, r, o;
+      unpack(batch[tasks[t].key].key, l, r, o);
+      partial[t] = count_exact(l, r, o, tasks[t].begin, tasks[t].end);
+    }
+    std::vector<int64_t> exact(batch.size(), 0);
+    for (size_t t = 0; t < tasks.size(); ++t) exact[tasks[t].key] += partial[t];
+    for (size_t k = 0; k < batch.size(); ++k) {
+      Stat& s = stats_[batch[k].key];
+      s.exact = exact[k];
+      s.exact_version = s.version;
+      s.added_at_exact = s.added;
+      if (s.exact >= min_count_) queue(batch[k].key, s, s.exact);
+    }
+  }
+
+  // Pairs are applied in root order; ranges split at safe frames are independent, so
+  // each is replayed sequentially by one thread with exactly the serial outcome.
+  void apply(int32_t left, int32_t right, int32_t off, int32_t child) {
+    struct Timer {
+      double& sink;
+      std::chrono::steady_clock::time_point t;
+      ~Timer() { sink += seconds_since(t); }
+    } timer{apply_seconds_, std::chrono::steady_clock::now()};
+    lefts(left);
+    lefts(right);
+    // Candidates are collected first: applying merges kills instances in both lists.
+    const auto collect_start = std::chrono::steady_clock::now();
+    std::vector<std::pair<InstanceId, InstanceId>> pairs = collect(left, right, off);
+    if (st_.tok.size() + pairs.size() > static_cast<size_t>(max_instances_)) {
+      pairs.clear();
+      pairs.shrink_to_fit();
+      compact_instances();
+      pairs = collect(left, right, off);
+      if (st_.tok.size() + pairs.size() > static_cast<size_t>(max_instances_))
+        throw std::length_error("live instance capacity exceeded after compaction");
+    }
+    collect_seconds_ += seconds_since(collect_start);
+    const std::vector<size_t> cuts = split(pairs);
+    const size_t groups = cuts.size() - 1;
+    const InstanceId base = st_.reserve(pairs.size());
+    if (pool_.size() < groups) pool_.resize(groups, std::vector<DeltaMap>(kShards));
+    dirty_shards_.resize(kShards);
+    std::vector<std::vector<InstanceId>> born(groups);
+    std::vector<int64_t> applied(groups, 0), rejected(groups, 0);
+    const auto loop_start = std::chrono::steady_clock::now();
+#ifdef _OPENMP
+#pragma omp parallel for schedule(dynamic, 1) num_threads(std::min<size_t>(threads_, groups))
+#endif
+    for (size_t g = 0; g < groups; ++g) {
+      auto& delta = pool_[g];
+      auto add = [&](uint64_t key, int64_t d) { delta[shard_of(key)].add(key, d); };
+      for (size_t k = cuts[g]; k < cuts[g + 1]; ++k) {
+        const InstanceId a = pairs[k].first, b = pairs[k].second;
+        if (!st_.alive[a] || !st_.alive[b]) continue;
+        if (!st_.valid(a, b, off)) {
+          ++rejected[g];
+          continue;
+        }
+        st_.partners(a, allow_zero_, [&](InstanceId x, InstanceId y, int32_t o) {
+          add(rule_key(st_.tok[x], st_.tok[y], o), -1);
+        });
+        st_.partners(b, allow_zero_, [&](InstanceId x, InstanceId y, int32_t o) {
+          if (x != a && y != a) add(rule_key(st_.tok[x], st_.tok[y], o), -1);
+        });
+        const int64_t r = st_.root[a];
+        st_.kill(a);
+        st_.kill(b);
+        const InstanceId m = base + static_cast<InstanceId>(k);
+        st_.tok[m] = child;
+        st_.root[m] = r;
+        st_.place(m);
+        born[g].push_back(m);
+        st_.partners(m, allow_zero_, [&](InstanceId x, InstanceId y, int32_t o) {
+          add(rule_key(st_.tok[x], st_.tok[y], o), +1);
+        });
+        ++applied[g];
+      }
+    }
+    (groups > 1 ? parallel_seconds_ : serial_seconds_) += seconds_since(loop_start);
+    const auto merge_start = std::chrono::steady_clock::now();
+    std::vector<int64_t> changes(kShards, 0);
+#ifdef _OPENMP
+#pragma omp parallel for schedule(dynamic, 1) num_threads(threads_) if (groups > 1)
+#endif
+    for (size_t p = 0; p < kShards; ++p) {
+      StatMap& table = stats_.shard(p);
+      for (size_t g = 0; g < groups; ++g) {
+        pool_[g][p].each([&](uint64_t k, int64_t d) {
+          change(k, d, table, dirty_shards_[p]);
+          ++changes[p];
+        });
+        pool_[g][p].clear();
+      }
+    }
+    for (size_t p = 0; p < kShards; ++p) {
+      dirty_.insert(dirty_.end(), dirty_shards_[p].begin(), dirty_shards_[p].end());
+      dirty_shards_[p].clear();
+      changes_ += changes[p];
+    }
+    for (size_t g = 0; g < groups; ++g) {
+      auto& occ = occurrence_[child];
+      occ.insert(occ.end(), born[g].begin(), born[g].end());
+      dead_[left] += applied[g];
+      dead_[right] += applied[g];
+      applied_ += applied[g];
+      rejected_ += rejected[g];
+      live_ -= applied[g];
+    }
+    merge_seconds_ += seconds_since(merge_start);
+  }
+
+  // Cut root-ordered pairs into contiguous groups at safe frames.
+  std::vector<size_t> split(const std::vector<std::pair<InstanceId, InstanceId>>& pairs) const {
+    std::vector<size_t> cuts{0};
+    const size_t n = pairs.size();
+    const size_t groups = std::max<size_t>(1, std::min<size_t>(threads_, n / kApplyGrain));
+    for (size_t g = 1; g < groups; ++g) {
+      const size_t target = std::max(cuts.back(), n * g / groups);
+      if (target >= n) break;
+      auto frame = std::upper_bound(safe_.begin(), safe_.end(), st_.root[pairs[target].first]);
+      if (frame == safe_.end()) break;
+      const size_t cut = static_cast<size_t>(
+          std::lower_bound(pairs.begin() + target, pairs.end(), *frame,
+                           [&](const std::pair<InstanceId, InstanceId>& p, int64_t f) {
+                             return st_.root[p.first] < f;
+                           }) -
+          pairs.begin());
+      if (cut > cuts.back() && cut < n) cuts.push_back(cut);
+    }
+    cuts.push_back(n);
+    return cuts;
+  }
+
+  static bool certified(const Stat& s) { return s.exact >= 0 && s.exact_version == s.version; }
+
+  static int64_t bound(const Stat& s) {
+    if (certified(s)) return s.exact;
+    return s.exact >= 0 ? std::min(s.raw, s.exact + s.added - s.added_at_exact) : s.raw;
+  }
+
+  void queue(uint64_t key, Stat& s, int64_t count) {
+    heap_.push({count, key, s.version});
+    s.live = true;
+    s.queued_version = s.version;
+    s.queued_count = count;
+  }
+
+  bool take(Entry& e, Stat*& s) {
+    e = heap_.top();
+    heap_.pop();
+    s = stats_.find(e.key);
+    if (!s || s->retired || !s->live || e.version != s->queued_version) return false;
+    s->live = false;
+    if (e.version == s->version) return true;
+    const int64_t b = bound(*s);
+    if (b > 0) queue(e.key, *s, b);
+    return false;
+  }
+
+  bool current_certified(const Entry& e) {
+    const Stat* s = stats_.find(e.key);
+    return s && !s->retired && s->live && e.version == s->queued_version &&
+           e.version == s->version && certified(*s);
+  }
+
+  void run(bool progress) {
+    live_ = initial_instances_;
+    while (tables_.size() < vocab_size_ && !heap_.empty()) {
+      Entry top;
+      Stat* found;
+      if (!take(top, found)) continue;
+      Stat& s = *found;
+      if (!certified(s)) {
+        std::vector<Entry> batch{top};
+        while (!heap_.empty() && batch.size() < kCertifyBatch) {
+          if (current_certified(heap_.top())) break;
+          Entry e;
+          Stat* next;
+          if (take(e, next)) batch.push_back(e);
+        }
+        certify(batch);
+        continue;
+      }
+      if (top.count != s.exact) continue;
+      if (s.exact < min_count_) continue;
+      int32_t l, r, o;
+      unpack(top.key, l, r, o);
+      std::vector<uint32_t> m;
+      if (!tables_.compose(l, r, o, m)) {
+        s.exact = 0;
+        continue;
+      }
+      const int32_t child = tables_.add_merge(l, r, o);
+      merges_.emplace_back(l, r, o);
+      s.retired = true;
+      dirty_.clear();
+      ++epoch_;
+      apply(l, r, o, child);
+      // Lazy decrease-key: only a grown bound needs a new entry now.
+      for (uint64_t k : dirty_) {
+        Stat& d = stats_[k];
+        if (d.retired) continue;
+        const int64_t b = bound(d);
+        if (b > 0 && (!d.live || b > d.queued_count)) queue(k, d, b);
+      }
+      dirty_.clear();
+      if (progress && merges_.size() % 256 == 0) {
+        py::gil_scoped_acquire acquire;
+        py::print("sparse temporal merges", merges_.size(), "live", live_, py::arg("flush") = true);
+      }
+    }
+  }
+
+  Tables tables_;
+  Instances st_;
+  int64_t vocab_size_;
+  bool allow_zero_;
+  bool dense_census_;
+  int64_t min_count_;
+  int64_t max_instances_;
+  std::vector<std::vector<InstanceId>> occurrence_;
+  ShardedStats stats_;
+  std::vector<std::vector<DeltaMap>> pool_;            // [apply group][shard]
+  std::vector<std::vector<uint64_t>> dirty_shards_;  // [shard]
+  std::priority_queue<Entry> heap_;
+  std::vector<uint64_t> dirty_;
+  std::vector<std::tuple<int32_t, int32_t, int32_t>> merges_;
+  std::vector<int64_t> dead_;
+  std::vector<int64_t> safe_;
+  uint32_t epoch_ = 0;
+  int threads_ = 1;
+  int64_t initial_instances_ = 0, live_ = 0, dry_runs_ = 0, rejected_ = 0, applied_ = 0;
+  int64_t compactions_ = 0;
+  size_t initial_types_ = 0;
+  int64_t changes_ = 0;
+  double census_seconds_ = 0, dry_seconds_ = 0, apply_seconds_ = 0, collect_seconds_ = 0,
+         parallel_seconds_ = 0, serial_seconds_ = 0, merge_seconds_ = 0;
+};
+
+struct Spatial {
+  std::vector<int64_t> offsets;
+  std::vector<uint32_t> support;
+  std::vector<std::vector<int32_t>> expansion;  // token -> value per sparse slot (-1 absent)
+  std::unordered_map<uint64_t, std::pair<int32_t, int32_t>> rules;  // (lo, hi) -> (child, rank)
+  int32_t dims = 0;
+
+  static uint64_t pair_key(int32_t a, int32_t b) {
+    if (a > b) std::swap(a, b);
+    return (static_cast<uint64_t>(a) << 32) | static_cast<uint32_t>(b);
+  }
+
+  void init(const std::vector<int64_t>& bins, const std::vector<std::pair<int32_t, int32_t>>& merges) {
+    dims = static_cast<int32_t>(bins.size());
+    offsets.assign(1, 0);
+    for (int64_t b : bins) {
+      if (b < 1) throw std::invalid_argument("every sparse slot needs at least one value");
+      offsets.push_back(offsets.back() + b);
+    }
+    for (int32_t d = 0; d < dims; ++d)
+      for (int64_t v = 0; v < bins[d]; ++v) {
+        support.push_back(1u << d);
+        std::vector<int32_t> e(dims, -1);
+        e[d] = static_cast<int32_t>(v);
+        expansion.push_back(std::move(e));
+      }
+    const int32_t atoms = static_cast<int32_t>(support.size());
+    for (size_t r = 0; r < merges.size(); ++r) {
+      const auto [a, b] = merges[r];
+      const int32_t child = atoms + static_cast<int32_t>(r);
+      if (a < 0 || b < 0 || a >= child || b >= child || (support[a] & support[b]))
+        throw std::invalid_argument("invalid spatial merge");
+      support.push_back(support[a] | support[b]);
+      std::vector<int32_t> e(dims);
+      for (int32_t d = 0; d < dims; ++d) e[d] = std::max(expansion[a][d], expansion[b][d]);
+      expansion.push_back(std::move(e));
+      rules.emplace(pair_key(a, b), std::make_pair(child, static_cast<int32_t>(r)));
+    }
+  }
+
+  std::vector<int32_t> encode(const int32_t* row) const {
+    std::set<int32_t> tokens;
+    for (int32_t d = 0; d < dims; ++d)
+      if (row[d] >= 0) tokens.insert(static_cast<int32_t>(offsets[d] + row[d]));
+    using Item = std::tuple<int32_t, int32_t, int32_t, int32_t>;
+    std::priority_queue<Item, std::vector<Item>, std::greater<Item>> heap;
+    auto offer = [&](int32_t a, int32_t b) {
+      auto it = rules.find(pair_key(a, b));
+      if (it != rules.end())
+        heap.emplace(it->second.second, it->second.first, std::min(a, b), std::max(a, b));
+    };
+    for (auto i = tokens.begin(); i != tokens.end(); ++i)
+      for (auto j = std::next(i); j != tokens.end(); ++j) offer(*i, *j);
+    while (!heap.empty()) {
+      const auto [rank, child, left, right] = heap.top();
+      heap.pop();
+      if (!tokens.count(left) || !tokens.count(right)) continue;
+      tokens.erase(left);
+      tokens.erase(right);
+      for (int32_t other : tokens) offer(child, other);
+      tokens.insert(child);
+    }
+    return {tokens.begin(), tokens.end()};
+  }
+};
+
+// Runtime codec over dense primitive cells ``(T, D)`` with ``-1`` for inactive slots.
+class SparseCodec {
+ public:
+  SparseCodec(std::vector<int64_t> bins, std::vector<int64_t> zero,
+              std::vector<std::pair<int32_t, int32_t>> spatial_merges, int32_t lookahead,
+              std::vector<std::tuple<int32_t, int32_t, int32_t>> temporal_merges)
+      : bins_(std::move(bins)), zero_(std::move(zero)) {
+    if (bins_.empty() || bins_.size() != zero_.size())
+      throw std::invalid_argument("bins and zero must have one entry per slot");
+    dims_ = static_cast<int32_t>(bins_.size());
+    std::vector<int64_t> sparse;
+    for (int32_t d = 0; d < dims_; ++d) {
+      // -2 retains every dense value; -1 change-codes an absolute slot.
+      if (zero_[d] >= bins_[d] || zero_[d] < -2) throw std::invalid_argument("zero index");
+      sparse.push_back(bins_[d] - (zero_[d] >= 0));
+    }
+    sparse.push_back(1);  // E
+    spatial_.init(sparse, spatial_merges);
+    tables_.init(spatial_.support, dims_ + 1, lookahead);
+    for (size_t r = 0; r < temporal_merges.size(); ++r) {
+      auto [a, b, o] = temporal_merges[r];
+      if (a < 0 || b < 0 || a >= tables_.size() || b >= tables_.size())
+        throw std::invalid_argument("merge order");
+      tables_.add_merge(a, b, o);
+      if (!rank_.emplace(rule_key(a, b, o), static_cast<int32_t>(r)).second)
+        throw std::invalid_argument("duplicate temporal rule");
+    }
+  }
+
+  int32_t vocab_size() const { return tables_.size(); }
+  int32_t spatial_vocab_size() const { return static_cast<int32_t>(spatial_.support.size()); }
+
+  std::vector<std::vector<int32_t>> encode_batch(
+      const py::array_t<int32_t, py::array::c_style | py::array::forcecast>& cells,
+      int threads) const {
+    if (cells.ndim() != 3 || cells.shape(2) != dims_ || cells.shape(0) < 1 || cells.shape(1) < 1)
+      throw std::invalid_argument("cells must have shape (B, T, D) with B, T >= 1");
+    const int64_t batch = cells.shape(0), steps = cells.shape(1);
+    const int32_t* data = cells.data();
+    std::vector<std::vector<int32_t>> out(batch);
+    {
+      py::gil_scoped_release release;
+      ac2::parallel_for(batch, threads, [&](int64_t b) {
+        out[b] = encode_one(data + b * steps * dims_, steps);
+      });
+    }
+    return out;
+  }
+
+  py::array_t<int32_t> decode_batch(const std::vector<std::vector<int32_t>>& rows,
+                                    int64_t horizon, const std::vector<bool>& active) const {
+    if (horizon < 1) throw std::invalid_argument("horizon must be positive");
+    if (static_cast<int32_t>(active.size()) != dims_)
+      throw std::invalid_argument("active must have one flag per slot");
+    py::array_t<int32_t> out({static_cast<int64_t>(rows.size()), horizon, int64_t{dims_}});
+    int32_t* data = out.mutable_data();
+    for (size_t k = 0; k < rows.size(); ++k) decode_one(rows[k], horizon, active, data + k * horizon * dims_);
+    return out;
+  }
+
+  // Spatial tokens of sparse frames ``(N, D + 1)``: per-frame counts and flat tokens.
+  py::tuple spatial_encode(
+      const py::array_t<int32_t, py::array::c_style | py::array::forcecast>& frames) const {
+    if (frames.ndim() != 2 || frames.shape(1) != dims_ + 1)
+      throw std::invalid_argument("frames must have shape (N, D + 1)");
+    const int64_t n = frames.shape(0);
+    py::array_t<int64_t> counts(n);
+    std::vector<int32_t> flat;
+    for (int64_t i = 0; i < n; ++i) {
+      const int32_t* row = frames.data() + i * (dims_ + 1);
+      for (int32_t d = 0; d <= dims_; ++d)
+        if (row[d] >= static_cast<int32_t>(spatial_.offsets[d + 1] - spatial_.offsets[d]))
+          throw std::invalid_argument("sparse value out of range");
+      const auto tokens = spatial_.encode(row);
+      counts.mutable_data()[i] = static_cast<int64_t>(tokens.size());
+      flat.insert(flat.end(), tokens.begin(), tokens.end());
+    }
+    return py::make_tuple(counts, py::array_t<int32_t>(flat.size(), flat.data()));
+  }
+
+  std::vector<std::vector<std::pair<int32_t, int32_t>>> patterns() const { return tables_.pattern; }
+
+ private:
+  // Dense chunk -> sparse frames (T, D + 1); validates slot activity and ranges.
+  std::vector<int32_t> sparsify(const int32_t* cells, int64_t steps) const {
+    const int32_t width = dims_ + 1;
+    std::vector<int32_t> sparse(static_cast<size_t>(steps) * width, -1);
+    int64_t last = -1;
+    for (int32_t d = 0; d < dims_; ++d) {
+      const bool active = cells[d] >= 0;
+      for (int64_t t = 0; t < steps; ++t) {
+        const int32_t v = cells[t * dims_ + d];
+        if ((v >= 0) != active)
+          throw std::invalid_argument("a slot must be active for the whole chunk or not at all");
+        if (v >= bins_[d]) throw std::invalid_argument("cell outside the primitive range");
+        if (!active) continue;
+        int32_t s = -1;
+        if (zero_[d] == -2) {
+          s = v;
+        } else if (zero_[d] < 0) {
+          if (t == 0 || v != cells[(t - 1) * dims_ + d]) s = v;
+        } else if (v != zero_[d]) {
+          s = v - (v > zero_[d]);
+        }
+        sparse[t * width + d] = s;
+        if (s >= 0) last = std::max(last, t);
+      }
+    }
+    for (int64_t t = 0; t < last; ++t) sparse[t * width + dims_] = 0;
+    return sparse;
+  }
+
+  std::vector<int32_t> encode_one(const int32_t* cells, int64_t steps) const {
+    const std::vector<int32_t> sparse = sparsify(cells, steps);
+    Instances s;
+    s.reset(&tables_, steps + tables_.lookahead);
+    for (int64_t t = 0; t < steps; ++t)
+      for (int32_t token : spatial_.encode(sparse.data() + t * (dims_ + 1))) s.add(token, t);
+    using Cand = std::tuple<int32_t, int64_t, int32_t, int64_t, InstanceId, InstanceId, int32_t>;
+    std::priority_queue<Cand, std::vector<Cand>, std::greater<Cand>> heap;
+    auto offer = [&](InstanceId i) {
+      s.partners(i, true, [&](InstanceId a, InstanceId b, int32_t off) {
+        auto it = rank_.find(rule_key(s.tok[a], s.tok[b], off));
+        if (it == rank_.end()) return;
+        heap.emplace(it->second, s.root[a], tables_.rdim[s.tok[a]], s.root[b], a, b, off);
+      });
+    };
+    for (size_t i = 0; i < s.tok.size(); ++i) offer(static_cast<InstanceId>(i));
+    while (!heap.empty()) {
+      auto [r, ra, da, rb, a, b, off] = heap.top();
+      heap.pop();
+      if (!s.alive[a] || !s.alive[b] || !s.valid(a, b, off)) continue;
+      const int64_t root = s.root[a];
+      s.kill(a);
+      s.kill(b);
+      offer(s.add(tables_.base + r, root));
+    }
+    return s.serialize();
+  }
+
+  void decode_one(const std::vector<int32_t>& row, int64_t horizon,
+                  const std::vector<bool>& active, int32_t* cells) const {
+    const int32_t L = tables_.lookahead, width = dims_ + 1;
+    uint32_t allowed = 1u << dims_;
+    for (int32_t d = 0; d < dims_; ++d)
+      if (active[d]) allowed |= 1u << d;
+    std::vector<uint32_t> cover(horizon + L, 0);
+    std::vector<int32_t> sparse(static_cast<size_t>(horizon) * width, -1);
+    int64_t front = 0;
+    std::vector<uint32_t> own(L, 0);
+    std::vector<int32_t> frames;
+    for (int32_t token : row) {
+      if (token < 0 || token >= tables_.size()) throw std::invalid_argument("token outside vocab");
+      if (front >= horizon) throw std::invalid_argument("token after the last open frame");
+      frames.clear();
+      for (auto [o, t] : tables_.pattern[token]) {
+        const int64_t u = front + o;
+        const uint32_t m = spatial_.support[t];
+        if (u >= horizon) throw std::invalid_argument("token exceeds horizon");
+        if (m & ~allowed) throw std::invalid_argument("token touches an inactive slot");
+        if (m & own[o]) throw std::invalid_argument("overlapping token supports");
+        if ((m >> dims_ & 1u) && u == horizon - 1)
+          throw std::invalid_argument("end-of-frame atom on the last frame");
+        if (!own[o]) frames.push_back(o);
+        own[o] |= m;
+        for (int32_t d = 0; d < width; ++d)
+          if (spatial_.expansion[t][d] >= 0) sparse[u * width + d] = spatial_.expansion[t][d];
+      }
+      for (int32_t o : frames) {
+        const int64_t u = front + o;
+        if (cover[u] & own[o]) throw std::invalid_argument("overlapping token supports");
+        if (cover[u] >> dims_ & 1u) throw std::invalid_argument("token touches a closed frame");
+        cover[u] |= own[o];
+        own[o] = 0;
+      }
+      while (front < horizon && (cover[front] >> dims_ & 1u)) ++front;
+    }
+    for (int32_t d = 0; d < dims_; ++d) {
+      int32_t carried = -1;
+      for (int64_t t = 0; t < horizon; ++t) {
+        const int32_t v = sparse[t * width + d];
+        int32_t& out = cells[t * dims_ + d];
+        if (!active[d]) {
+          out = -1;
+        } else if (zero_[d] == -2) {
+          if (v < 0) throw std::invalid_argument("dense active slot is missing");
+          out = v;
+        } else if (zero_[d] < 0) {
+          if (v >= 0) carried = v;
+          if (carried < 0) throw std::invalid_argument("an absolute slot must be stated on frame 0");
+          out = carried;
+        } else {
+          out = v < 0 ? static_cast<int32_t>(zero_[d]) : v + (v >= zero_[d]);
+        }
+      }
+    }
+  }
+
+  std::vector<int64_t> bins_, zero_;
+  int32_t dims_ = 0;
+  Spatial spatial_;
+  Tables tables_;
+  std::unordered_map<uint64_t, int32_t> rank_;
+};
+
+}  // namespace
+
+void register_sparse(py::module_& m) {
+  py::class_<Trainer>(m, "SparseTrainer")
+      .def(py::init<std::vector<uint32_t>, int32_t, int32_t, int64_t, bool, int64_t, int, bool, int64_t>(),
+           py::arg("spatial_support"), py::arg("dims"), py::arg("lookahead"),
+           py::arg("vocab_size"), py::arg("allow_zero_offset"), py::arg("min_count") = 2,
+           py::arg("threads") = 0, py::arg("dense_census") = false,
+           py::arg("max_instances") = UINT32_MAX)
+      .def("set_instances", &Trainer::set_instances)
+      .def("train", &Trainer::train, py::arg("progress") = false)
+      .def("merges", &Trainer::merges);
+  py::class_<SparseCodec>(m, "SparseCodec")
+      .def(py::init<std::vector<int64_t>, std::vector<int64_t>,
+                    std::vector<std::pair<int32_t, int32_t>>, int32_t,
+                    std::vector<std::tuple<int32_t, int32_t, int32_t>>>(),
+           py::arg("bins"), py::arg("zero"), py::arg("spatial_merges"), py::arg("lookahead"),
+           py::arg("temporal_merges"))
+      .def_property_readonly("vocab_size", &SparseCodec::vocab_size)
+      .def_property_readonly("spatial_vocab_size", &SparseCodec::spatial_vocab_size)
+      .def("encode_batch", &SparseCodec::encode_batch, py::arg("cells"), py::arg("threads") = 0)
+      .def("decode_batch", &SparseCodec::decode_batch, py::arg("rows"), py::arg("horizon"),
+           py::arg("active"))
+      .def("spatial_encode", &SparseCodec::spatial_encode, py::arg("frames"))
+      .def("patterns", &SparseCodec::patterns);
+}

actioncodec2_kernels/spatial.cpp

@@ -0,0 +1,297 @@
+
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <pybind11/stl.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+#ifdef _OPENMP
+#include <omp.h>
+#endif
+
+#include "common.hpp"
+
+namespace py = pybind11;
+
+namespace {
+
+inline uint64_t pair_key(int32_t left, int32_t right) {
+  uint32_t low = static_cast<uint32_t>(left < right ? left : right);
+  uint32_t high = static_cast<uint32_t>(left < right ? right : left);
+  return (static_cast<uint64_t>(low) << 32) | high;
+}
+
+struct Corpus {
+  int64_t records = 0;
+  int32_t width = 0;                 // slots per record == action dimension
+  std::vector<int32_t> member;       // [records * width], -1 when free
+  std::vector<int64_t> weights;      // [records]
+};
+
+}  // namespace
+
+py::tuple train_spatial_native(py::array_t<int32_t, py::array::c_style | py::array::forcecast> codes,
+                               py::array_t<int64_t, py::array::c_style | py::array::forcecast> weights,
+                               const std::vector<int32_t>& bins, int64_t vocab_size,
+                               int threads) {
+  auto codes_view = codes.unchecked<2>();
+  auto weights_view = weights.unchecked<1>();
+  const int64_t records = codes_view.shape(0);
+  const int32_t dimension = static_cast<int32_t>(codes_view.shape(1));
+  if (dimension != static_cast<int32_t>(bins.size())) {
+    throw std::invalid_argument("codes must have one column per bin count");
+  }
+  if (weights_view.shape(0) != records) {
+    throw std::invalid_argument("weights must have one entry per record");
+  }
+
+  std::vector<int32_t> offsets(dimension + 1, 0);
+  for (int32_t index = 0; index < dimension; ++index) offsets[index + 1] = offsets[index] + bins[index];
+  const int32_t atom_count = offsets[dimension];
+  if (vocab_size < atom_count) {
+    throw std::invalid_argument("vocab_size is below the number of primitive atoms");
+  }
+
+#ifdef _OPENMP
+  const int thread_count = threads > 0 ? threads : omp_get_max_threads();
+#else
+  (void)threads;
+  const int thread_count = 1;
+#endif
+
+  Corpus corpus;
+  corpus.records = records;
+  corpus.width = dimension;
+  corpus.member.resize(static_cast<size_t>(records) * dimension);
+  corpus.weights.resize(records);
+  int64_t total_weight = 0;
+  int64_t initial_length = 0;
+  for (int64_t record = 0; record < records; ++record) {
+    corpus.weights[record] = weights_view(record);
+    total_weight += weights_view(record);
+    int32_t active = 0;
+    for (int32_t slot = 0; slot < dimension; ++slot) {
+      const int32_t value = codes_view(record, slot);
+      if (value < -1 || value >= bins[slot]) {
+        throw std::invalid_argument("a cell index is outside its coordinate's bin count");
+      }
+      if (value < 0) {
+        corpus.member[static_cast<size_t>(record) * dimension + slot] = -1;
+      } else {
+        corpus.member[static_cast<size_t>(record) * dimension + slot] = offsets[slot] + value;
+        ++active;
+      }
+    }
+    if (active == 0) throw std::invalid_argument("every record needs an active coordinate");
+    initial_length += weights_view(record) * active;
+  }
+
+  // Occurrence sets, kept sorted so intersections and differences are linear.
+  std::vector<std::vector<int32_t>> occurrence(atom_count);
+  for (int32_t slot = 0; slot < dimension; ++slot) {
+    for (int64_t record = 0; record < records; ++record) {
+      const int32_t token = corpus.member[static_cast<size_t>(record) * dimension + slot];
+      if (token >= 0) occurrence[token].push_back(static_cast<int32_t>(record));
+    }
+  }
+  for (auto& members : occurrence) std::sort(members.begin(), members.end());
+
+  // --- initial census -------------------------------------------------------
+  ac2::CounterMap counts(static_cast<size_t>(records) * 4 + 64);
+  {
+    std::vector<ac2::CounterMap> shards;
+    shards.reserve(thread_count);
+    for (int index = 0; index < thread_count; ++index) shards.emplace_back(1024);
+#ifdef _OPENMP
+#pragma omp parallel num_threads(thread_count)
+#endif
+    {
+#ifdef _OPENMP
+      const int thread = omp_get_thread_num();
+#else
+      const int thread = 0;
+#endif
+      ac2::CounterMap& local = shards[thread];
+#ifdef _OPENMP
+#pragma omp for schedule(static)
+#endif
+      for (int64_t record = 0; record < records; ++record) {
+        const int32_t* row = &corpus.member[static_cast<size_t>(record) * dimension];
+        const int64_t weight = corpus.weights[record];
+        for (int32_t a = 0; a < dimension; ++a) {
+          if (row[a] < 0) continue;
+          for (int32_t b = a + 1; b < dimension; ++b) {
+            if (row[b] < 0) continue;
+            local.add(pair_key(row[a], row[b]), weight);
+          }
+        }
+      }
+    }
+    for (auto& shard : shards) {
+      shard.for_each([&](uint64_t key, int64_t value) { counts.add(key, value); });
+    }
+  }
+
+  std::vector<ac2::HeapEntry> heap;
+  heap.reserve(counts.size());
+  counts.for_each([&](uint64_t key, int64_t value) { heap.push_back({value, key, 0}); });
+  std::make_heap(heap.begin(), heap.end(), ac2::HeapOrder{});
+  const size_t initial_pair_types = heap.size();
+
+  // --- greedy merging -------------------------------------------------------
+  std::vector<std::pair<int32_t, int32_t>> merges;
+  merges.reserve(static_cast<size_t>(vocab_size - atom_count));
+  int64_t length = initial_length;  // total active tokens over the corpus
+  int64_t heap_pops = 0;
+  int64_t counter_updates = 0;
+
+  std::vector<int32_t> matched;
+  std::vector<int32_t> remainder;
+  std::vector<int32_t> affected;
+  std::vector<int32_t> stamp(static_cast<size_t>(vocab_size) + 16, 0);
+  int32_t generation = 0;
+  std::vector<ac2::CounterMap> deltas;
+  deltas.reserve(thread_count);
+  for (int index = 0; index < thread_count; ++index) deltas.emplace_back(1024);
+
+  while (static_cast<int64_t>(merges.size()) + atom_count < vocab_size) {
+    int32_t left = -1;
+    int32_t right = -1;
+    int64_t frequency = 0;
+    while (!heap.empty()) {
+      const ac2::HeapEntry top = ac2::heap_pop(heap);
+      ++heap_pops;
+      const int64_t current = counts.get(top.key);
+      if (current <= 0) continue;
+      if (current != top.count) continue;  // a fresher entry for this key exists
+      left = static_cast<int32_t>(top.key >> 32);
+      right = static_cast<int32_t>(top.key & 0xffffffffu);
+      frequency = current;
+      break;
+    }
+    if (left < 0) break;  // every record has collapsed to a single token
+
+    matched.clear();
+    std::set_intersection(occurrence[left].begin(), occurrence[left].end(),
+                          occurrence[right].begin(), occurrence[right].end(),
+                          std::back_inserter(matched));
+    if (matched.empty()) {  // cannot happen while counters are exact
+      counts.add(pair_key(left, right), -frequency);
+      continue;
+    }
+
+    const int32_t child = static_cast<int32_t>(occurrence.size());
+    occurrence.emplace_back();
+
+    const bool parallel = matched.size() > 4096 && thread_count > 1;
+    if (parallel) {
+      for (auto& shard : deltas) shard = ac2::CounterMap(1024);
+    }
+#ifdef _OPENMP
+#pragma omp parallel if (parallel) num_threads(thread_count)
+#endif
+    {
+#ifdef _OPENMP
+      const int thread = parallel ? omp_get_thread_num() : 0;
+#else
+      const int thread = 0;
+#endif
+      ac2::CounterMap& local = parallel ? deltas[thread] : counts;
+#ifdef _OPENMP
+#pragma omp for schedule(static)
+#endif
+      for (int64_t index = 0; index < static_cast<int64_t>(matched.size()); ++index) {
+        const int32_t record = matched[index];
+        const int32_t* row = &corpus.member[static_cast<size_t>(record) * dimension];
+        const int64_t weight = corpus.weights[record];
+        for (int32_t slot = 0; slot < dimension; ++slot) {
+          const int32_t other = row[slot];
+          if (other < 0 || other == left || other == right) continue;
+          local.add(pair_key(left, other), -weight);
+          local.add(pair_key(right, other), -weight);
+          local.add(pair_key(child, other), weight);
+        }
+      }
+    }
+    if (parallel) {
+      for (auto& shard : deltas) {
+        shard.for_each([&](uint64_t key, int64_t value) { counts.add(key, value); });
+      }
+    }
+
+    // Rewrite the touched rows: the child takes the left parent's slot.
+    for (const int32_t record : matched) {
+      int32_t* row = &corpus.member[static_cast<size_t>(record) * dimension];
+      for (int32_t slot = 0; slot < dimension; ++slot) {
+        if (row[slot] == left) {
+          row[slot] = child;
+        } else if (row[slot] == right) {
+          row[slot] = -1;
+        }
+      }
+    }
+
+    counts.add(pair_key(left, right), -frequency);
+    counter_updates += static_cast<int64_t>(matched.size()) * dimension;
+
+    occurrence[child] = matched;
+    for (const int32_t parent : {left, right}) {
+      remainder.clear();
+      std::set_difference(occurrence[parent].begin(), occurrence[parent].end(), matched.begin(),
+                          matched.end(), std::back_inserter(remainder));
+      occurrence[parent].swap(remainder);
+    }
+
+    affected.clear();
+    if (static_cast<size_t>(child) >= stamp.size()) stamp.resize(child + 1024, 0);
+    ++generation;
+    for (const int32_t record : matched) {
+      const int32_t* row = &corpus.member[static_cast<size_t>(record) * dimension];
+      for (int32_t slot = 0; slot < dimension; ++slot) {
+        const int32_t other = row[slot];
+        if (other < 0 || other == child) continue;
+        if (stamp[other] == generation) continue;
+        stamp[other] = generation;
+        affected.push_back(other);
+      }
+    }
+    for (const int32_t other : affected) {
+      for (const int32_t parent : {left, right, child}) {
+        const uint64_t key = pair_key(parent, other);
+        const int64_t value = counts.get(key);
+        if (value > 0) ac2::heap_push(heap, {value, key, 0});
+      }
+    }
+
+    merges.emplace_back(left, right);
+    length -= frequency;
+
+    if (heap.size() > 4 * counts.size() + 65536) {
+      heap.clear();
+      counts.for_each([&](uint64_t key, int64_t value) { heap.push_back({value, key, 0}); });
+      std::make_heap(heap.begin(), heap.end(), ac2::HeapOrder{});
+    }
+  }
+
+  py::list merge_list;
+  for (const auto& merge : merges) {
+    merge_list.append(py::make_tuple(merge.first, merge.second));
+  }
+  py::dict report;
+  report["records"] = records;
+  report["weight"] = total_weight;
+  report["atoms"] = atom_count;
+  report["merges"] = static_cast<int64_t>(merges.size());
+  report["vocab_size"] = atom_count + static_cast<int64_t>(merges.size());
+  report["tokens_per_action_initial"] = static_cast<float>(initial_length) / total_weight;
+  report["tokens_per_action"] = static_cast<float>(length) / static_cast<float>(total_weight);
+  report["compression"] = static_cast<float>(initial_length) / static_cast<float>(length);
+  report["initial_pair_types"] = static_cast<int64_t>(initial_pair_types);
+  report["heap_pops"] = heap_pops;
+  report["counter_updates"] = counter_updates;
+  report["threads"] = thread_count;
+  report["backend"] = "native_sparse_exact";
+  return py::make_tuple(merge_list, report);
+}

actioncodec2_kernels/temporal.cpp

@@ -0,0 +1,799 @@
+// Native exact rooted temporal Set-BPE trainer.
+//
+// A temporal symbol is rooted at one token of the current control step and may
+// absorb spatial tokens from a bounded number of future steps.  Training is the
+// exact sparse local-graph algorithm: counters exist only for rule keys that
+// actually occur inside the lookahead window, each merge updates only the edges
+// incident to the instances it changed, and shifted self-rules (left == right)
+// are certified by exact left-to-right maximum matching before they may win.
+//
+// This mirrors actioncodec2/setbpe/temporal.py merge for merge, including the
+// tie-break, so the reference trainer can prove the two agree.  What differs is
+// the constant factor: the Python trainer allocates a dense
+// ``lookahead x V x V`` legality table, while this one derives legality from the
+// per-token time-shifted support masks in O(span) with no allocation, which is
+// what makes large vocabularies tractable.
+
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <pybind11/stl.h>
+
+#include <algorithm>
+#include <array>
+#include <chrono>
+#include <cstdint>
+#include <deque>
+#include <limits>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "common.hpp"
+
+#ifdef _OPENMP
+#include <omp.h>
+#endif
+
+namespace py = pybind11;
+
+namespace {
+
+struct TemporalEntry {
+  int64_t count;
+  uint64_t key;
+  uint32_t version;
+  uint8_t certified;
+};
+
+struct TemporalOrder {
+  bool operator()(const TemporalEntry& a, const TemporalEntry& b) const {
+    if (a.count != b.count) return a.count < b.count;        // larger count first
+    if (a.key != b.key) return a.key > b.key;                // smaller key first
+    return a.certified > b.certified;                        // uncertified first
+  }
+};
+
+struct InstanceChunk {
+  const int32_t* tokens;
+  const int64_t* anchors;
+  int64_t size;
+};
+
+size_t initial_counter_capacity(int64_t instance_count, int32_t base_vocab, int32_t lookahead) {
+  const size_t by_instances = static_cast<size_t>(instance_count) * 4 + 64;
+  const size_t vocab = static_cast<size_t>(base_vocab);
+  const size_t offsets = static_cast<size_t>(std::max(lookahead - 1, 0));
+  if (vocab != 0 && offsets > (std::numeric_limits<size_t>::max() - 64) / vocab / vocab) {
+    return by_instances;
+  }
+  const size_t by_key_universe = vocab * vocab * offsets + 64;
+  return std::min(by_instances, by_key_universe);
+}
+
+inline uint64_t rule_key(int32_t left, int32_t right, int32_t offset) {
+  return (static_cast<uint64_t>(left) << 40) | (static_cast<uint64_t>(right) << 16) |
+         static_cast<uint64_t>(offset);
+}
+
+inline void unpack(uint64_t key, int32_t& left, int32_t& right, int32_t& offset) {
+  left = static_cast<int32_t>(key >> 40);
+  right = static_cast<int32_t>((key >> 16) & 0xffffffull);
+  offset = static_cast<int32_t>(key & 0xffffull);
+}
+
+class Trainer {
+ public:
+  Trainer(int64_t total_anchors, int64_t instance_capacity_hint,
+          const uint64_t* spatial_support, const int32_t* spatial_anchor, int32_t base_vocab,
+          int32_t dimension, int32_t lookahead, int64_t vocab_size, int threads)
+      : base_vocab_(base_vocab),
+        dimension_(dimension),
+        lookahead_(lookahead),
+        total_anchors_(total_anchors),
+        counts_(initial_counter_capacity(instance_capacity_hint, base_vocab, lookahead)) {
+#ifdef _OPENMP
+    // More workers eventually make the per-merge local-map reduction slower
+    // than the incident scan itself.  Keep auto mode at the measured scaling
+    // knee; an explicit value remains an opt-in override.
+    threads_ = threads > 0 ? threads : std::min(omp_get_max_threads(), 32);
+#else
+    (void)threads;
+    threads_ = 1;
+#endif
+    const size_t capacity = static_cast<size_t>(vocab_size) + 16;
+    masks_.assign(capacity * lookahead_, 0);
+    spans_.assign(capacity, 1);
+    root_anchor_.assign(capacity, 0);
+    pattern_start_.assign(capacity + 1, 0);
+    for (int32_t token = 0; token < base_vocab_; ++token) {
+      masks_[static_cast<size_t>(token) * lookahead_] = spatial_support[token];
+      root_anchor_[token] = spatial_anchor[token];
+      pattern_start_[token] = static_cast<int32_t>(pattern_.size());
+      pattern_.push_back({0, token});
+    }
+    pattern_start_[base_vocab_] = static_cast<int32_t>(pattern_.size());
+    token_count_ = base_vocab_;
+
+    member_.assign(static_cast<size_t>(total_anchors) * dimension_, -1);
+    occurrence_.resize(capacity);
+    dirty_.assign(capacity, 0);
+  }
+
+  void add_chunk(const InstanceChunk& chunk) {
+    for (int64_t index = 0; index < chunk.size; ++index) {
+      add_instance(chunk.tokens[index], chunk.anchors[index]);
+    }
+    initial_count_ += chunk.size;
+  }
+
+  void finalize_instances() {
+    if (initial_count_ <= 0) throw std::runtime_error("temporal training needs instances");
+    active_count_ = initial_count_;
+    const auto census_started = std::chrono::steady_clock::now();
+    census();
+    census_seconds_ = elapsed(census_started);
+    counts_.for_each([&](uint64_t key, int64_t value) { heap_.push_back({value, key, 0, 0}); });
+    std::make_heap(heap_.begin(), heap_.end(), TemporalOrder{});
+    initial_edge_types_ = heap_.size();
+  }
+
+  // -- geometry ------------------------------------------------------------
+
+  bool legal(int32_t left, int32_t right, int32_t offset) const {
+    const int32_t right_span = spans_[right];
+    if (offset + right_span > lookahead_) return false;
+    const uint64_t* left_masks = &masks_[static_cast<size_t>(left) * lookahead_];
+    const uint64_t* right_masks = &masks_[static_cast<size_t>(right) * lookahead_];
+    const int32_t overlap = std::min(spans_[left] - offset, right_span);
+    for (int32_t step = 0; step < overlap; ++step) {
+      if (left_masks[offset + step] & right_masks[step]) return false;
+    }
+    return true;
+  }
+
+  // -- instance bookkeeping ------------------------------------------------
+
+  int32_t add_instance(int32_t token, int64_t anchor) {
+    if (token < 0 || static_cast<size_t>(token) >= occurrence_.size()) {
+      throw std::runtime_error("temporal instance token " + std::to_string(token) +
+                               " is outside trainer capacity " +
+                               std::to_string(occurrence_.size()));
+    }
+    if (anchor < 0 || anchor >= total_anchors_) {
+      throw std::runtime_error("temporal instance anchor is outside the corpus timeline");
+    }
+    int32_t* row = &member_[static_cast<size_t>(anchor) * dimension_];
+    int32_t slot = 0;
+    while (slot < dimension_ && row[slot] >= 0) ++slot;
+    if (slot == dimension_) {
+      throw std::runtime_error("temporal anchor exceeded the token-width bound");
+    }
+    row[slot] = token;
+    occurrence_[token].push_back(anchor);
+    dirty_[token] = 1;
+    return slot;
+  }
+
+  // Removal touches only the row mirror.  The occurrence list is allowed to
+  // keep stale anchors because ``occurs_at`` is an exact membership test, so
+  // compaction can be deferred to the next time the list is actually read.
+  void remove_instance(int32_t token, int64_t anchor) {
+    int32_t* row = &member_[static_cast<size_t>(anchor) * dimension_];
+    for (int32_t slot = 0; slot < dimension_; ++slot) {
+      if (row[slot] == token) {
+        row[slot] = -1;
+        break;
+      }
+    }
+    dirty_[token] = 1;
+  }
+
+  const std::vector<int64_t>& sorted_occurrence(int32_t token) {
+    auto& anchors = occurrence_[token];
+    if (dirty_[token]) {
+      std::sort(anchors.begin(), anchors.end());
+      anchors.erase(std::unique(anchors.begin(), anchors.end()), anchors.end());
+      anchors.erase(std::remove_if(anchors.begin(), anchors.end(),
+                                   [&](int64_t anchor) { return !occurs_at(token, anchor); }),
+                    anchors.end());
+      dirty_[token] = 0;
+    }
+    return anchors;
+  }
+
+  bool occurs_at(int32_t token, int64_t anchor) const {
+    const int32_t* row = &member_[static_cast<size_t>(anchor) * dimension_];
+    for (int32_t slot = 0; slot < dimension_; ++slot) {
+      if (row[slot] == token) return true;
+    }
+    return false;
+  }
+
+  // -- matching ------------------------------------------------------------
+
+  std::vector<std::pair<int64_t, int64_t>> select(int32_t left, int32_t right, int32_t offset) {
+    std::vector<std::pair<int64_t, int64_t>> matches;
+    const std::vector<int64_t>& anchors = sorted_occurrence(left);
+    if (left != right) {
+#ifdef _OPENMP
+      if (threads_ > 1 && anchors.size() >= 65536) {
+        const int workers = std::min<int64_t>(threads_, anchors.size() / 32768);
+        std::vector<std::vector<std::pair<int64_t, int64_t>>> local(workers);
+#pragma omp parallel num_threads(workers)
+        {
+          const int worker = omp_get_thread_num();
+          const size_t start = anchors.size() * static_cast<size_t>(worker) / workers;
+          const size_t stop = anchors.size() * static_cast<size_t>(worker + 1) / workers;
+          auto& output = local[worker];
+          output.reserve(stop - start);
+          for (size_t index = start; index < stop; ++index) {
+            const int64_t anchor = anchors[index];
+            const int64_t target = anchor + offset;
+            if (target < total_anchors_ && occurs_at(right, target)) {
+              output.emplace_back(anchor, target);
+            }
+          }
+        }
+        size_t count = 0;
+        for (const auto& output : local) count += output.size();
+        matches.reserve(count);
+        for (auto& output : local) {
+          matches.insert(matches.end(), output.begin(), output.end());
+        }
+        return matches;
+      }
+#endif
+      for (const int64_t anchor : anchors) {
+        const int64_t target = anchor + offset;
+        if (target < total_anchors_ && occurs_at(right, target)) {
+          matches.emplace_back(anchor, target);
+        }
+      }
+      return matches;
+    }
+    // Shifted self-rule: an instance cannot be both the left and the right end
+    // of two different merges, so the true gain is a maximum matching, which
+    // for a fixed shift on a line is exactly greedy left to right.
+    std::deque<int64_t> consumed_targets;
+    for (const int64_t anchor : anchors) {
+      while (!consumed_targets.empty() && consumed_targets.front() < anchor) {
+        consumed_targets.pop_front();
+      }
+      if (!consumed_targets.empty() && consumed_targets.front() == anchor) {
+        consumed_targets.pop_front();
+        continue;
+      }
+      const int64_t wide = anchor + offset;
+      if (wide >= total_anchors_) continue;
+      const int64_t target = wide;
+      if (!occurs_at(left, target)) continue;
+      matches.emplace_back(anchor, target);
+      consumed_targets.push_back(target);
+    }
+    return matches;
+  }
+
+  // -- search --------------------------------------------------------------
+
+  bool best(int32_t& left, int32_t& right, int32_t& offset,
+            std::vector<std::pair<int64_t, int64_t>>& matches) {
+    while (!heap_.empty()) {
+      std::pop_heap(heap_.begin(), heap_.end(), TemporalOrder{});
+      const TemporalEntry top = heap_.back();
+      heap_.pop_back();
+      ++heap_pops_;
+      const int64_t current = counts_.get(top.key);
+      if (current <= 0) continue;
+      unpack(top.key, left, right, offset);
+      const bool self_rule = left == right;
+      const uint32_t version = self_rule ? self_version(top.key) : 0;
+      const bool stale = self_rule ? version != top.version : current != top.count;
+      if (stale) {
+        heap_.push_back({current, top.key, version, 0});
+        std::push_heap(heap_.begin(), heap_.end(), TemporalOrder{});
+        continue;
+      }
+      if (self_rule && !top.certified) {
+        // Raw edge count only upper-bounds the achievable gain here; replace
+        // the priority with the exact matching size and re-queue.
+        auto certified = select(left, right, offset);
+        self_cache_[top.key] = {version, certified};
+        heap_.push_back({static_cast<int64_t>(certified.size()), top.key, version, 1});
+        std::push_heap(heap_.begin(), heap_.end(), TemporalOrder{});
+        ++self_certifications_;
+        continue;
+      }
+      if (top.certified) {
+        auto cached = self_cache_.find(top.key);
+        if (cached == self_cache_.end() || cached->second.first != version) continue;
+        matches = cached->second.second;
+      } else {
+        matches = select(left, right, offset);
+      }
+      if (static_cast<int64_t>(matches.size()) != top.count) {
+        throw std::runtime_error("temporal counter disagrees with exact matching");
+      }
+      return !matches.empty();
+    }
+    return false;
+  }
+
+  // -- commit --------------------------------------------------------------
+
+  void commit(int32_t left, int32_t right, int32_t offset,
+              const std::vector<std::pair<int64_t, int64_t>>& matches) {
+    const int32_t child = token_count_++;
+    // The child's pattern is the left pattern plus the right pattern shifted.
+    const int32_t start = static_cast<int32_t>(pattern_.size());
+    for (int32_t index = pattern_start_[left]; index < pattern_start_[left + 1]; ++index) {
+      pattern_.push_back(pattern_[index]);
+    }
+    for (int32_t index = pattern_start_[right]; index < pattern_start_[right + 1]; ++index) {
+      pattern_.push_back({pattern_[index].first + offset, pattern_[index].second});
+    }
+    std::sort(pattern_.begin() + start, pattern_.end());
+    pattern_start_[child] = start;
+    pattern_start_[child + 1] = static_cast<int32_t>(pattern_.size());
+
+    int32_t span = 0;
+    int32_t roots = 0;
+    uint64_t* child_masks = &masks_[static_cast<size_t>(child) * lookahead_];
+    for (int32_t index = start; index < static_cast<int32_t>(pattern_.size()); ++index) {
+      const int32_t time = pattern_[index].first;
+      const uint64_t support = masks_[static_cast<size_t>(pattern_[index].second) * lookahead_];
+      if (time >= lookahead_) throw std::runtime_error("temporal merge exceeds the lookahead");
+      if (child_masks[time] & support) {
+        throw std::runtime_error("temporal merge joins overlapping spatial supports");
+      }
+      child_masks[time] |= support;
+      span = std::max(span, time + 1);
+      roots += (time == 0);
+    }
+    if (roots != 1) throw std::runtime_error("a rooted temporal token needs one time-zero token");
+    spans_[child] = span;
+    root_anchor_[child] = root_anchor_[left];
+
+    ac2::CounterMap deltas(
+        std::min(matches.size() * 8 + 64, incident_key_capacity(token_count_)));
+    std::vector<std::pair<int32_t, int64_t>> removed;  // (token, anchor), sorted by anchor
+    removed.reserve(matches.size() * 2);
+    size_t left_index = 0;
+    size_t right_index = 0;
+    while (left_index < matches.size() || right_index < matches.size()) {
+      const int64_t left_anchor =
+          left_index < matches.size() ? matches[left_index].first : total_anchors_;
+      const int64_t right_anchor =
+          right_index < matches.size() ? matches[right_index].second : total_anchors_;
+      if (left_anchor <= right_anchor) {
+        removed.emplace_back(left, left_anchor);
+        ++left_index;
+      } else {
+        removed.emplace_back(right, right_anchor);
+        ++right_index;
+      }
+    }
+    accumulate_incident(removed, -1, deltas);
+    correct_internal(removed, +1, deltas);
+
+    for (const auto& instance : removed) remove_instance(instance.first, instance.second);
+    std::vector<std::pair<int32_t, int64_t>> children;
+    children.reserve(matches.size());
+    for (const auto& match : matches) {
+      add_instance(child, match.first);
+      children.emplace_back(child, match.first);
+    }
+    accumulate_incident(children, +1, deltas);
+    correct_internal(children, -1, deltas);
+
+    apply(deltas, rule_key(left, right, offset));
+    active_count_ -= static_cast<int64_t>(matches.size());
+    merges_.push_back({left, right, offset});
+  }
+
+  // -- accessors -----------------------------------------------------------
+
+  const std::vector<std::array<int32_t, 3>>& merges() const { return merges_; }
+  int64_t active_count() const { return active_count_; }
+  int64_t initial_count() const { return initial_count_; }
+  int64_t heap_pops() const { return heap_pops_; }
+  int64_t counter_updates() const { return counter_updates_; }
+  int64_t self_certifications() const { return self_certifications_; }
+  size_t initial_edge_types() const { return initial_edge_types_; }
+  size_t live_edge_types() const { return counts_.live_entries(); }
+  size_t counter_buckets() const { return counts_.buckets(); }
+  int threads() const { return threads_; }
+  double census_seconds() const { return census_seconds_; }
+  double incident_seconds() const { return incident_seconds_; }
+  double incident_parallel_seconds() const { return incident_parallel_seconds_; }
+  double incident_merge_seconds() const { return incident_merge_seconds_; }
+  double correction_seconds() const { return correction_seconds_; }
+  double correction_parallel_seconds() const { return correction_parallel_seconds_; }
+  double apply_seconds() const { return apply_seconds_; }
+
+ private:
+  static double elapsed(std::chrono::steady_clock::time_point started) {
+    return std::chrono::duration<double>(std::chrono::steady_clock::now() - started).count();
+  }
+
+  size_t incident_key_capacity(int32_t vocabulary) const {
+    return static_cast<size_t>(6) * static_cast<size_t>(vocabulary) *
+               static_cast<size_t>(std::max(lookahead_ - 1, 0)) +
+           64;
+  }
+
+  void census_anchor(int64_t anchor, ac2::CounterMap& counter) const {
+    const int32_t* row = &member_[static_cast<size_t>(anchor) * dimension_];
+    for (int32_t slot = 0; slot < dimension_; ++slot) {
+      const int32_t left = row[slot];
+      if (left < 0) continue;
+      for (int32_t offset = 1; offset < lookahead_; ++offset) {
+        const int64_t target = anchor + offset;
+        if (target >= total_anchors_) break;
+        const int32_t* other_row = &member_[static_cast<size_t>(target) * dimension_];
+        for (int32_t other_slot = 0; other_slot < dimension_; ++other_slot) {
+          const int32_t right = other_row[other_slot];
+          if (right >= 0 && legal(left, right, offset)) {
+            counter.add(rule_key(left, right, offset), 1);
+          }
+        }
+      }
+    }
+  }
+
+  void census() {
+    const size_t key_universe = initial_counter_capacity(total_anchors_, base_vocab_, lookahead_);
+#ifdef _OPENMP
+    if (threads_ > 1 && total_anchors_ >= 4096) {
+      const int workers = std::min<int64_t>(threads_, total_anchors_ / 2048);
+      // Small rule universes fit comfortably in one local map per worker.
+      // For a large spatial vocabulary, each long-running worker would
+      // eventually discover almost the whole universe and multiply memory by
+      // the thread count.  Merge fixed-size anchor batches instead: this keeps
+      // local memory bounded while the global counter remains exact.
+      const bool batched = key_universe > (1u << 20);
+      const int64_t anchors_per_worker = batched ? 8192 : total_anchors_;
+      const int64_t batch_size = anchors_per_worker * workers;
+      const size_t local_hint = batched
+                                    ? std::min<size_t>(key_universe,
+                                                       anchors_per_worker * 32 + 64)
+                                    : key_universe;
+      std::vector<std::unique_ptr<ac2::CounterMap>> local;
+      local.reserve(workers);
+      for (int worker = 0; worker < workers; ++worker) {
+        local.push_back(std::make_unique<ac2::CounterMap>(local_hint));
+      }
+      for (int64_t batch_start = 0; batch_start < total_anchors_; batch_start += batch_size) {
+        const int64_t batch_end = std::min(total_anchors_, batch_start + batch_size);
+#pragma omp parallel for schedule(static) num_threads(workers)
+        for (int64_t anchor = batch_start; anchor < batch_end; ++anchor) {
+          census_anchor(anchor, *local[omp_get_thread_num()]);
+        }
+        for (const auto& counter : local) {
+          counter->for_each([&](uint64_t key, int64_t value) { counts_.add(key, value); });
+          if (batched) counter->clear();
+        }
+      }
+      return;
+    }
+#endif
+    for (int64_t anchor = 0; anchor < total_anchors_; ++anchor) census_anchor(anchor, counts_);
+  }
+
+  uint32_t self_version(uint64_t key) const {
+    auto found = self_versions_.find(key);
+    return found == self_versions_.end() ? 0u : found->second;
+  }
+
+  void accumulate_incident(const std::vector<std::pair<int32_t, int64_t>>& instances,
+                           int64_t multiplier, ac2::CounterMap& deltas) {
+    const auto started = std::chrono::steady_clock::now();
+#ifdef _OPENMP
+    if (threads_ > 1 && instances.size() >= 1024) {
+      const int workers = std::min<int>(threads_, static_cast<int>(instances.size() / 512));
+      const size_t hint = std::min(
+          instances.size() * 8 / static_cast<size_t>(workers) + 64,
+          incident_key_capacity(token_count_));
+      std::vector<std::unique_ptr<ac2::CounterMap>> local;
+      local.reserve(workers);
+      for (int worker = 0; worker < workers; ++worker) {
+        local.push_back(std::make_unique<ac2::CounterMap>(hint));
+      }
+      const auto parallel_started = std::chrono::steady_clock::now();
+#pragma omp parallel for schedule(static) num_threads(workers)
+      for (int64_t index = 0; index < static_cast<int64_t>(instances.size()); ++index) {
+        accumulate_one(instances[static_cast<size_t>(index)], multiplier,
+                       *local[omp_get_thread_num()]);
+      }
+      incident_parallel_seconds_ += elapsed(parallel_started);
+      const auto merge_started = std::chrono::steady_clock::now();
+      for (const auto& counter : local) {
+        counter->for_each([&](uint64_t key, int64_t value) { deltas.add(key, value); });
+      }
+      incident_merge_seconds_ += elapsed(merge_started);
+      incident_seconds_ += elapsed(started);
+      return;
+    }
+#endif
+    for (const auto& instance : instances) {
+      accumulate_one(instance, multiplier, deltas);
+    }
+    incident_seconds_ += elapsed(started);
+  }
+
+  void accumulate_one(const std::pair<int32_t, int64_t>& instance, int64_t multiplier,
+                      ac2::CounterMap& deltas) const {
+    const int32_t token = instance.first;
+    const int64_t anchor = instance.second;
+    for (int32_t offset = 1; offset < lookahead_; ++offset) {
+      const int64_t forward = anchor + offset;
+      if (forward < total_anchors_) {
+        const int32_t* row = &member_[static_cast<size_t>(forward) * dimension_];
+        for (int32_t slot = 0; slot < dimension_; ++slot) {
+          const int32_t other = row[slot];
+          if (other >= 0 && legal(token, other, offset)) {
+            deltas.add(rule_key(token, other, offset), multiplier);
+          }
+        }
+      }
+      const int64_t backward = anchor - offset;
+      if (backward >= 0) {
+        const int32_t* row = &member_[static_cast<size_t>(backward) * dimension_];
+        for (int32_t slot = 0; slot < dimension_; ++slot) {
+          const int32_t other = row[slot];
+          if (other >= 0 && legal(other, token, offset)) {
+            deltas.add(rule_key(other, token, offset), multiplier);
+          }
+        }
+      }
+    }
+  }
+
+  // Incident aggregation sees a changed-to-changed edge once from each end.
+  void correct_internal(const std::vector<std::pair<int32_t, int64_t>>& instances,
+                        int64_t correction, ac2::CounterMap& deltas) {
+    const auto started = std::chrono::steady_clock::now();
+    if (instances.empty()) return;
+    std::vector<size_t> groups{0};
+    groups.reserve(instances.size());
+    for (size_t index = 1; index < instances.size(); ++index) {
+      if (instances[index].second < instances[index - 1].second) {
+        throw std::runtime_error("changed temporal instances must be anchor-sorted");
+      }
+      if (instances[index].second != instances[index - 1].second) groups.push_back(index);
+    }
+    groups.push_back(instances.size());
+    const size_t group_count = groups.size() - 1;
+#ifdef _OPENMP
+    if (threads_ > 1 && group_count >= 1024) {
+      const int workers = std::min<int>(threads_, static_cast<int>(group_count / 512));
+      const size_t hint = std::min(group_count * 2 / static_cast<size_t>(workers) + 64,
+                                   incident_key_capacity(token_count_));
+      std::vector<std::unique_ptr<ac2::CounterMap>> local;
+      local.reserve(workers);
+      for (int worker = 0; worker < workers; ++worker) {
+        local.push_back(std::make_unique<ac2::CounterMap>(hint));
+      }
+      const auto parallel_started = std::chrono::steady_clock::now();
+#pragma omp parallel for schedule(static) num_threads(workers)
+      for (int64_t index = 0; index < static_cast<int64_t>(group_count); ++index) {
+        correct_group(static_cast<size_t>(index), groups, instances, correction,
+                      *local[omp_get_thread_num()]);
+      }
+      correction_parallel_seconds_ += elapsed(parallel_started);
+      for (const auto& counter : local) {
+        counter->for_each([&](uint64_t key, int64_t value) { deltas.add(key, value); });
+      }
+    } else
+#endif
+    {
+      for (size_t group = 0; group < group_count; ++group) {
+        correct_group(group, groups, instances, correction, deltas);
+      }
+    }
+    correction_seconds_ += elapsed(started);
+  }
+
+  void correct_group(size_t group, const std::vector<size_t>& groups,
+                     const std::vector<std::pair<int32_t, int64_t>>& instances,
+                     int64_t correction, ac2::CounterMap& deltas) const {
+    const size_t left_start = groups[group];
+    const size_t left_stop = groups[group + 1];
+    const int64_t anchor = instances[left_start].second;
+    for (size_t other = group + 1; other + 1 < groups.size(); ++other) {
+      const size_t right_start = groups[other];
+      const int64_t distance = instances[right_start].second - anchor;
+      if (distance >= lookahead_) break;
+      const int32_t offset = static_cast<int32_t>(distance);
+      for (size_t left_index = left_start; left_index < left_stop; ++left_index) {
+        const int32_t left = instances[left_index].first;
+        for (size_t right_index = right_start; right_index < groups[other + 1]; ++right_index) {
+          const int32_t right = instances[right_index].first;
+          if (legal(left, right, offset)) {
+            deltas.add(rule_key(left, right, offset), correction);
+          }
+        }
+      }
+    }
+  }
+
+  void apply(const ac2::CounterMap& deltas, uint64_t selected) {
+    const auto started = std::chrono::steady_clock::now();
+    deltas.for_each([&](uint64_t key, int64_t delta) {
+      const auto [previous, count] = counts_.add_with_previous(key, delta);
+      if (count < 0) {
+        int32_t bl = 0, br = 0, bo = 0;
+        unpack(key, bl, br, bo);
+        throw std::runtime_error("temporal edge counter became negative: rule=(" +
+                                 std::to_string(bl) + "," + std::to_string(br) + "," +
+                                 std::to_string(bo) + ") previous=" + std::to_string(previous) +
+                                 " delta=" + std::to_string(delta) + " at merge " +
+                                 std::to_string(merges_.size()));
+      }
+      counter_updates_ += delta < 0 ? -delta : delta;
+      int32_t left = 0;
+      int32_t right = 0;
+      int32_t offset = 0;
+      unpack(key, left, right, offset);
+      uint32_t version = 0;
+      if (left == right) {
+        version = self_version(key) + 1;
+        self_versions_[key] = version;
+        self_cache_.erase(key);
+      }
+      // Every edge a commit adds contains the freshly minted child token, so a
+      // pre-existing key can only lose mass; its stale heap priority stays an
+      // admissible upper bound and is repaired when popped.  New keys and the
+      // key just consumed need an entry right away.
+      if (count > 0 && (previous == 0 || key == selected)) {
+        heap_.push_back({count, key, version, 0});
+        std::push_heap(heap_.begin(), heap_.end(), TemporalOrder{});
+      }
+    });
+    apply_seconds_ += elapsed(started);
+  }
+
+  int32_t base_vocab_;
+  int32_t dimension_;
+  int32_t lookahead_;
+  int64_t total_anchors_;
+  int threads_ = 1;
+  int32_t token_count_ = 0;
+
+  std::vector<uint64_t> masks_;
+  std::vector<int32_t> spans_;
+  std::vector<int32_t> root_anchor_;
+  std::vector<std::pair<int32_t, int32_t>> pattern_;  // (time, spatial token)
+  std::vector<int32_t> pattern_start_;
+
+  std::vector<int32_t> member_;
+  std::vector<std::vector<int64_t>> occurrence_;
+  std::vector<uint8_t> dirty_;
+
+  ac2::CounterMap counts_;
+  std::vector<TemporalEntry> heap_;
+  std::unordered_map<uint64_t, uint32_t> self_versions_;
+  std::unordered_map<uint64_t, std::pair<uint32_t, std::vector<std::pair<int64_t, int64_t>>>>
+      self_cache_;
+
+  std::vector<std::array<int32_t, 3>> merges_;
+  int64_t active_count_ = 0;
+  int64_t initial_count_ = 0;
+  int64_t heap_pops_ = 0;
+  int64_t counter_updates_ = 0;
+  int64_t self_certifications_ = 0;
+  size_t initial_edge_types_ = 0;
+  double census_seconds_ = 0.0;
+  double incident_seconds_ = 0.0;
+  double incident_parallel_seconds_ = 0.0;
+  double incident_merge_seconds_ = 0.0;
+  double correction_seconds_ = 0.0;
+  double correction_parallel_seconds_ = 0.0;
+  double apply_seconds_ = 0.0;
+};
+
+}  // namespace
+
+py::tuple train_temporal_native(
+    py::iterable instance_chunks, int64_t total_anchors,
+    py::array_t<uint64_t, py::array::c_style | py::array::forcecast> spatial_support,
+    py::array_t<int32_t, py::array::c_style | py::array::forcecast> spatial_anchor,
+    int32_t dimension, int32_t lookahead, int64_t vocab_size, int64_t action_count, int threads) {
+  if (lookahead < 1) throw std::invalid_argument("lookahead must be positive");
+  const int32_t base_vocab = static_cast<int32_t>(spatial_support.size());
+  if (vocab_size < base_vocab) {
+    throw std::invalid_argument("vocab_size is below the spatial vocabulary size");
+  }
+  std::unique_ptr<Trainer> trainer;
+  {
+    py::gil_scoped_release release;
+    trainer = std::make_unique<Trainer>(
+        total_anchors, action_count, spatial_support.data(),
+        spatial_anchor.data(), base_vocab, dimension, lookahead, vocab_size, threads);
+  }
+  int64_t instance_count = 0;
+  for (py::handle raw_chunk : instance_chunks) {
+    py::tuple chunk = py::cast<py::tuple>(raw_chunk);
+    if (chunk.size() != 2) {
+      throw std::invalid_argument("every temporal instance chunk must be a token/anchor pair");
+    }
+    auto tokens = py::cast<
+        py::array_t<int32_t, py::array::c_style | py::array::forcecast>>(chunk[0]);
+    auto anchors = py::cast<
+        py::array_t<int64_t, py::array::c_style | py::array::forcecast>>(chunk[1]);
+    if (tokens.ndim() != 1 || anchors.ndim() != 1 || tokens.size() != anchors.size()) {
+      throw std::invalid_argument("every temporal token/anchor chunk must be an aligned vector");
+    }
+    for (py::ssize_t index = 0; index < tokens.size(); ++index) {
+      if (tokens.data()[index] < 0 || tokens.data()[index] >= base_vocab) {
+        throw std::invalid_argument("input temporal token is outside the spatial vocabulary");
+      }
+    }
+    const InstanceChunk view{tokens.data(), anchors.data(), tokens.size()};
+    {
+      py::gil_scoped_release release;
+      trainer->add_chunk(view);
+    }
+    instance_count += tokens.size();
+  }
+  if (instance_count <= 0) {
+    throw std::invalid_argument("temporal instance chunks must not be empty");
+  }
+  int32_t left = 0;
+  int32_t right = 0;
+  int32_t offset = 0;
+  std::vector<std::pair<int64_t, int64_t>> matches;
+  double selection_seconds = 0.0;
+  double commit_seconds = 0.0;
+  {
+    py::gil_scoped_release release;
+    trainer->finalize_instances();
+    while (static_cast<int64_t>(trainer->merges().size()) + base_vocab < vocab_size) {
+      auto started = std::chrono::steady_clock::now();
+      const bool found = trainer->best(left, right, offset, matches);
+      selection_seconds +=
+          std::chrono::duration<double>(std::chrono::steady_clock::now() - started).count();
+      if (!found) break;
+      started = std::chrono::steady_clock::now();
+      trainer->commit(left, right, offset, matches);
+      commit_seconds +=
+          std::chrono::duration<double>(std::chrono::steady_clock::now() - started).count();
+    }
+  }
+
+  py::list merge_list;
+  for (const auto& merge : trainer->merges()) {
+    merge_list.append(py::make_tuple(merge[0], merge[1], merge[2]));
+  }
+  py::dict report;
+  report["lookahead"] = lookahead;
+  report["spatial_vocab_size"] = base_vocab;
+  report["merges"] = static_cast<int64_t>(trainer->merges().size());
+  report["vocab_size"] = base_vocab + static_cast<int64_t>(trainer->merges().size());
+  report["actions"] = action_count;
+  report["tokens_per_action_initial"] =
+      static_cast<float>(trainer->initial_count()) / static_cast<float>(action_count);
+  report["tokens_per_action"] =
+      static_cast<float>(trainer->active_count()) / static_cast<float>(action_count);
+  report["compression"] =
+      static_cast<float>(trainer->initial_count()) / static_cast<float>(trainer->active_count());
+  report["initial_edge_types"] = static_cast<int64_t>(trainer->initial_edge_types());
+  report["live_edge_types"] = static_cast<int64_t>(trainer->live_edge_types());
+  report["counter_buckets"] = static_cast<int64_t>(trainer->counter_buckets());
+  report["census_seconds"] = trainer->census_seconds();
+  report["incident_seconds"] = trainer->incident_seconds();
+  report["incident_parallel_seconds"] = trainer->incident_parallel_seconds();
+  report["incident_merge_seconds"] = trainer->incident_merge_seconds();
+  report["correction_seconds"] = trainer->correction_seconds();
+  report["correction_parallel_seconds"] = trainer->correction_parallel_seconds();
+  report["apply_seconds"] = trainer->apply_seconds();
+  report["selection_seconds"] = selection_seconds;
+  report["commit_seconds"] = commit_seconds;
+  report["heap_pops"] = trainer->heap_pops();
+  report["counter_updates"] = trainer->counter_updates();
+  report["self_certifications"] = trainer->self_certifications();
+  report["backend"] = "native_sparse_exact";
+  report["threads"] = trainer->threads();
+  report["threads_requested"] = threads;
+  return py::make_tuple(merge_list, report);
+}

actioncodec2_kernels/token_streams.hpp

@@ -0,0 +1,48 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+#include <cstdint>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace ac2 {
+namespace py = pybind11;
+
+// Preserve the public Python integer contract; never coerce floats or booleans.
+inline std::vector<std::vector<int64_t>> checked_token_streams(
+    const py::iterable& streams, int64_t vocabulary, bool routed = false) {
+  const py::object numpy_integer = py::module_::import("numpy").attr("integer");
+  py::object integral;
+  std::vector<std::vector<int64_t>> output;
+  if (PyList_Check(streams.ptr())) output.reserve(PyList_GET_SIZE(streams.ptr()));
+  for (py::handle stream : streams) {
+    std::vector<int64_t> row;
+    if (PyList_Check(stream.ptr())) row.reserve(PyList_GET_SIZE(stream.ptr()));
+    for (py::handle raw : py::reinterpret_borrow<py::iterable>(stream)) {
+      bool integer = PyLong_Check(raw.ptr()) || py::isinstance(raw, numpy_integer);
+      if (!integer && !routed) {
+        if (!integral) integral = py::module_::import("numbers").attr("Integral");
+        integer = py::isinstance(raw, integral);
+      }
+      if (PyBool_Check(raw.ptr()) || !integer)
+        throw std::invalid_argument(routed ? "tokens[" + std::to_string(row.size()) + "] must be an integer"
+                                           : "tokens must be integers");
+      py::object number = py::reinterpret_steal<py::object>(PyNumber_Long(raw.ptr()));
+      if (!number) throw py::error_already_set();
+      int overflow = 0;
+      const int64_t token = PyLong_AsLongLongAndOverflow(number.ptr(), &overflow);
+      if (PyErr_Occurred()) throw py::error_already_set();
+      if (overflow || token < 0 || token >= vocabulary)
+        throw std::invalid_argument(routed ? "tokens[" + std::to_string(row.size()) +
+                                           "] is outside the routed profile vocabularies; out of range [0, " + std::to_string(vocabulary) + ")"
+                                           : "token is outside the vocabulary");
+      row.push_back(token);
+    }
+    if (row.empty()) throw std::invalid_argument("token stream must not be empty");
+    output.push_back(std::move(row));
+  }
+  if (output.empty()) throw std::invalid_argument("streams must not be empty");
+  return output;
+}
+}  // namespace ac2

config.json

@@ -1,4 +1,7 @@
 {
   "model_type": "actioncodec2",
-  "processor_class": "ActionCodec2"
+  "processor_class": "ActionCodec2",
+  "auto_map": {
+    "AutoProcessor": "processing_actioncodec2.ActionCodec2"
+  }
 }

processing_actioncodec2.py

@@ -21,11 +21,15 @@ SHARED_PATTERNS = (
 
 def asset_inventory(directory: Path) -> dict[str, str]:
     """Hash the complete configuration and fitted assets for one variant."""
-    names = ["router_config.yaml", "config.json", "processor_config.json"]
+    sparse = (directory / "sparse_router_config.yaml").is_file()
+    names = ["sparse_router_config.yaml" if sparse else "router_config.yaml", "config.json", "processor_config.json"]
     profiles = [name for name in ("eef", "joint") if (directory / "profiles" / name).is_dir()]
     if not profiles:
         raise ValueError("variant has no fitted profiles")
-    names.extend(f"profiles/{profile}/{name}" for profile in profiles for name in PROFILE_FILES)
+    names.extend(f"profiles/{profile}/{name}" for profile in profiles
+                 for name in (tuple("sparse_bpe.json" if n == "bpe.json" else n for n in PROFILE_FILES) if sparse else PROFILE_FILES))
+    names.extend(f"profiles/{profile}/preprocessing.json" for profile in profiles
+                 if (directory / "profiles" / profile / "preprocessing.json").is_file())
     return {name: hashlib.sha256((directory / name).read_bytes()).hexdigest()
             for name in sorted(names)}
 
@@ -62,11 +66,12 @@ class _VariantBinding:
 
     def save_pretrained(self, save_directory, push_to_hub=False, **kwargs):
         directory = Path(save_directory)
+        include_training_report = kwargs.pop("include_training_report", False)
         if (directory / "variants.json").is_file():
             raise ValueError("save a selected variant to a separate directory, outside the catalog root")
         if not push_to_hub and kwargs:
             raise TypeError(f"unused save_pretrained keyword(s): {', '.join(sorted(kwargs))}")
-        files = super().save_pretrained(directory)
+        files = super().save_pretrained(directory, include_training_report=include_training_report)
         inventory = asset_inventory(directory)
         manifest = directory / "artifact_manifest.json"
         manifest.write_text(json.dumps({
@@ -77,10 +82,8 @@ class _VariantBinding:
         }, indent=2) + "\n", encoding="utf-8")
         files.append(str(manifest))
         if push_to_hub:
-            package = type(self.action_tokenizer).__module__.split(".routing.", 1)[0]
-            importlib.import_module(f"{package}.integration.hf").upload_pretrained_directory(
-                directory, **kwargs
-            )
+            package = type(self).__module__.split(".frontend.", 1)[0]
+            importlib.import_module(f"{package}.backends.actioncodec2.hf").upload_pretrained_directory(directory, **kwargs)
         return files
 
 
@@ -169,7 +172,7 @@ class ActionCodec2:
                 root = Path(snapshot_download(
                     str(pretrained_model_name_or_path),
                     allow_patterns=[*SHARED_PATTERNS, "router_config.yaml", "profiles/**",
-                                    "artifact_manifest.json"], **kwargs,
+                                    "sparse_router_config.yaml", "artifact_manifest.json"], **kwargs,
                 ))
             directory = root
             manifest = root / "artifact_manifest.json"
@@ -185,33 +188,28 @@ class ActionCodec2:
         if variant_fingerprint is not None and variant_fingerprint != origin.get("fingerprint"):
             raise ValueError("variant fingerprint differs from the tokenizer bound to this policy")
 
-        runtime = root / "runtime"
-        package_file = runtime / "__init__.py"
-        if not package_file.is_file():
-            raise FileNotFoundError(f"ActionCodec2 runtime is missing from {root}")
-        identity = hashlib.sha256(str(runtime.resolve()).encode()).hexdigest()[:16]
-        package_name = f"_actioncodec2_artifact_{identity}"
-        if package_name not in sys.modules:
-            spec = importlib.util.spec_from_file_location(
-                package_name, package_file, submodule_search_locations=[str(runtime)]
-            )
-            if spec is None or spec.loader is None:
-                raise ImportError(f"cannot load ActionCodec2 runtime from {runtime}")
-            package = importlib.util.module_from_spec(spec)
-            sys.modules[package_name] = package
-            try:
+        family = "actioncodec2_sparse" if (directory / "sparse_router_config.yaml").is_file() else "actioncodec2"
+        if __package__ == "actioncodec2.core":
+            package_name = __package__
+        else:
+            runtime = root / "runtime"
+            package_file = runtime / "__init__.py"
+            if not package_file.is_file():
+                raise FileNotFoundError(f"ActionCodec2 runtime is missing from {root}")
+            identity = hashlib.sha256(str(runtime.resolve()).encode()).hexdigest()[:16]
+            package_name = f"_actioncodec2_artifact_{identity}"
+            if package_name not in sys.modules:
+                spec = importlib.util.spec_from_file_location(
+                    package_name, package_file, submodule_search_locations=[str(runtime)]
+                )
+                package = importlib.util.module_from_spec(spec)
+                sys.modules[package_name] = package
                 spec.loader.exec_module(package)
-            except BaseException:
-                del sys.modules[package_name]
-                raise
-        processor = importlib.import_module(f"{package_name}.processing_actioncodec2")
-        if (runtime / "hub_acceleration.py").is_file():
-            acceleration = importlib.import_module(f"{package_name}.hub_acceleration")
-            acceleration.install(package_name)
+        processor = importlib.import_module(f"{package_name}.frontend.{family}.processor")
+        base = processor.ActionCodec2Sparse if family == "actioncodec2_sparse" else processor.ActionCodec2
         if not hasattr(processor, "_catalog_processor_class"):
             processor._catalog_processor_class = type(
-                "ActionCodec2", (_VariantBinding, processor.ActionCodec2),
-                {"__module__": processor.__name__},
+                base.__name__, (_VariantBinding, base), {"__module__": processor.__name__}
             )
         codec = processor._catalog_processor_class.from_pretrained(directory, action_space=action_space)
         codec._variant_origin = origin

pyproject.toml

@@ -1,16 +1,3 @@
 [build-system]
 requires = ["setuptools>=77", "wheel", "pybind11>=2.12"]
 build-backend = "setuptools.build_meta"
-
-[project]
-name = "actioncodec2-kernels"
-version = "0.1.0"
-description = "Portable C++17 encode/decode kernels for ActionCodec2 tokenizer artifacts"
-requires-python = ">=3.10"
-license = "Apache-2.0"
-
-[tool.setuptools]
-packages = ["actioncodec2_kernels"]
-
-[tool.setuptools.package-data]
-actioncodec2_kernels = ["*.cpp", "*.hpp"]

requirements.txt

@@ -1,6 +0,0 @@
-numpy
-scipy
-torch>=2.3
-transformers>=4.57,<6
-huggingface-hub
-pyyaml

runtime/__init__.py

@@ -1 +1 @@
-"""Private inference runtime bundled with this ActionCodec2 artifact."""
+"""Shared ActionCodec2 and ActionCodec2Sparse runtime."""

runtime/backend/__init__.py

@@ -1 +0,0 @@
-"""Internal backend components for ActionCodec2."""

runtime/backends/actioncodec2/__init__.py

@@ -0,0 +1,8 @@
+"""Shared physical models, layouts, and dense and sparse routers."""
+from .configuration import ActionCodec2Config, ConfigError
+from .presets import ComponentSpec, PresetConfig
+from .processor import ActionCodec2Processor
+from .routing import ActionCodec2Router
+from .trajectory import ActiveTrajectory
+
+__all__ = ["ActionCodec2Config", "ConfigError", "ComponentSpec", "PresetConfig", "ActionCodec2Processor", "ActionCodec2Router", "ActiveTrajectory"]

runtime/{physical → backends/actioncodec2}/_so3.py

@@ -4,9 +4,22 @@ from __future__ import annotations
 
 import numpy as np
 
+from .native_backend import native_kernel
+
+
+def _mapped(values, operation, width):
+    kernel = native_kernel("quaternion_map")
+    if kernel is None:
+        return None
+    array = np.asarray(values, dtype=np.float64)
+    return kernel(array.reshape(-1, array.shape[-1]), operation).reshape(*array.shape[:-1], width)
+
 
 def normalize_quaternion(values: np.ndarray) -> np.ndarray:
     """Normalize quaternion rows and reject an undefined zero rotation."""
+    result = _mapped(values, "normalize", 4)
+    if result is not None:
+        return result
     quaternions = np.asarray(values, dtype=np.float64)
     norms = np.linalg.norm(quaternions, axis=-1, keepdims=True)
     if np.any(norms < 1e-12):
@@ -18,6 +31,10 @@ def multiply_quaternions(left: np.ndarray, right: np.ndarray) -> np.ndarray:
     """Return the Hamilton product of broadcastable xyzw arrays."""
     first = np.asarray(left, dtype=np.float64)
     second = np.asarray(right, dtype=np.float64)
+    kernel = native_kernel("quaternion_multiply")
+    if kernel is not None:
+        first, second = np.broadcast_arrays(first, second)
+        return kernel(first.reshape(-1, 4), second.reshape(-1, 4)).reshape(first.shape)
     first_vector, first_scalar = first[..., :3], first[..., 3:4]
     second_vector, second_scalar = second[..., :3], second[..., 3:4]
     return normalize_quaternion(
@@ -36,6 +53,9 @@ def multiply_quaternions(left: np.ndarray, right: np.ndarray) -> np.ndarray:
 
 def conjugate_quaternion(quaternion: np.ndarray) -> np.ndarray:
     """Return the conjugate of an xyzw quaternion array."""
+    native = _mapped(quaternion, "conjugate", 4)
+    if native is not None:
+        return native
     result = np.asarray(quaternion, dtype=np.float64).copy()
     result[..., :3] *= -1.0
     return result
@@ -43,20 +63,24 @@ def conjugate_quaternion(quaternion: np.ndarray) -> np.ndarray:
 
 def rotvec_to_quaternion(rotvec: np.ndarray) -> np.ndarray:
     """Convert rotation vectors to xyzw quaternions."""
+    result = _mapped(rotvec, "from_rotvec", 4)
+    if result is not None:
+        return result
     values = np.asarray(rotvec, dtype=np.float64)
     angle = np.linalg.norm(values, axis=-1)
     scale = np.full_like(angle, 0.5)
     regular = angle >= 1e-12
     scale[regular] = np.sin(0.5 * angle[regular]) / angle[regular]
     return normalize_quaternion(
-        np.concatenate(
-            (values * scale[..., None], np.cos(0.5 * angle)[..., None]), axis=-1
-        )
+        np.concatenate((values * scale[..., None], np.cos(0.5 * angle)[..., None]), axis=-1)
     )
 
 
 def quaternion_to_rotvec(quaternion: np.ndarray) -> np.ndarray:
     """Convert xyzw quaternions to shortest-path rotation vectors."""
+    result = _mapped(quaternion, "to_rotvec", 3)
+    if result is not None:
+        return result
     values = normalize_quaternion(quaternion).copy()
     values[values[..., 3] < 0.0] *= -1.0
     vector_norm = np.linalg.norm(values[..., :3], axis=-1)
@@ -72,6 +96,9 @@ def euler_xyz_to_quaternion(euler: np.ndarray) -> np.ndarray:
     values = np.asarray(euler, dtype=np.float64)
     if values.ndim != 2 or values.shape[1] != 3:
         raise ValueError("XYZ Euler angles must have shape [N, 3]")
+    result = _mapped(values, "from_euler", 4)
+    if result is not None:
+        return result
     half = 0.5 * values
     sin_roll, sin_pitch, sin_yaw = np.sin(half).T
     cos_roll, cos_pitch, cos_yaw = np.cos(half).T
@@ -88,6 +115,9 @@ def quaternion_to_euler_xyz(quaternion: np.ndarray) -> np.ndarray:
     values = normalize_quaternion(quaternion)
     if values.ndim != 2 or values.shape[1] != 4:
         raise ValueError("xyzw quaternions must have shape [N, 4]")
+    result = _mapped(values, "to_euler", 3)
+    if result is not None:
+        return result
     x, y, z, w = values.T
     return np.stack(
         (

runtime/{common → backends/actioncodec2}/_validation.py

@@ -18,13 +18,11 @@ def active_dimensions(
 ) -> tuple[int, ...]:
     """Normalize and validate a non-empty increasing subset of dimensions."""
     dims = (
-        tuple(range(dimension))
-        if active_dims is None
-        else tuple(int(item) for item in active_dims)
+        tuple(range(dimension)) if active_dims is None else tuple(int(item) for item in active_dims)
     )
     if not dims:
         raise ValueError(f"{subject} must not be empty")
-    if any(left >= right for left, right in zip(dims, dims[1:])):
+    if any(left >= right for left, right in zip(dims, dims[1:], strict=False)):
         raise ValueError(f"{subject} must be strictly increasing")
     if dims[0] < 0 or dims[-1] >= dimension:
         raise ValueError(f"{subject} contains a dimension outside the vocabulary")
@@ -76,13 +74,9 @@ def validate_dimension_partition(
     for name, dims in groups:
         for index in dims:
             if index < 0 or index >= dimension:
-                raise error_type(
-                    f"{label} {index} used by {name!r} is outside input_dim"
-                )
+                raise error_type(f"{label} {index} used by {name!r} is outside input_dim")
             if index in owners:
-                raise error_type(
-                    f"{label} {index} belongs to both {owners[index]!r} and {name!r}"
-                )
+                raise error_type(f"{label} {index} belongs to both {owners[index]!r} and {name!r}")
             owners[index] = name
     missing = sorted(set(range(dimension)) - set(owners))
     if missing:

runtime/backends/actioncodec2/assets/presets/dual_eef_abs.yaml

@@ -0,0 +1,20 @@
+name: dual_eef_abs
+input_dim: 14
+input: absolute
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  left_arm:
+    space: eef
+    position_dims: [0, 1, 2]
+    rotation_dims: [3, 4, 5]
+    rotation_format: rotation_vector
+    gripper_dim: 6
+  right_arm:
+    space: eef
+    position_dims: [7, 8, 9]
+    rotation_dims: [10, 11, 12]
+    rotation_format: rotation_vector
+    gripper_dim: 13

runtime/backends/actioncodec2/assets/presets/dual_eef_delta.yaml

@@ -0,0 +1,20 @@
+name: dual_eef_delta
+input_dim: 14
+input: delta
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  left_arm:
+    space: eef
+    position_dims: [0, 1, 2]
+    rotation_dims: [3, 4, 5]
+    rotation_format: rotation_vector
+    gripper_dim: 6
+  right_arm:
+    space: eef
+    position_dims: [7, 8, 9]
+    rotation_dims: [10, 11, 12]
+    rotation_format: rotation_vector
+    gripper_dim: 13

runtime/backends/actioncodec2/assets/presets/dual_joint6_abs.yaml

@@ -0,0 +1,18 @@
+name: dual_joint6_abs
+input_dim: 14
+input: absolute
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  left_arm:
+    space: joint
+    joint_dims: [0, 1, 2, 3, 4, 5]
+    joint_wrap: pi
+    gripper_dim: 6
+  right_arm:
+    space: joint
+    joint_dims: [7, 8, 9, 10, 11, 12]
+    joint_wrap: pi
+    gripper_dim: 13

runtime/backends/actioncodec2/assets/presets/dual_joint6_delta.yaml

@@ -0,0 +1,16 @@
+name: dual_joint6_delta
+input_dim: 14
+input: delta
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  left_arm:
+    space: joint
+    joint_dims: [0, 1, 2, 3, 4, 5]
+    gripper_dim: 6
+  right_arm:
+    space: joint
+    joint_dims: [7, 8, 9, 10, 11, 12]
+    gripper_dim: 13

runtime/backends/actioncodec2/assets/presets/dual_joint7_abs.yaml

@@ -0,0 +1,18 @@
+name: dual_joint7_abs
+input_dim: 16
+input: absolute
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  left_arm:
+    space: joint
+    joint_dims: [0, 1, 2, 3, 4, 5, 6]
+    joint_wrap: pi
+    gripper_dim: 7
+  right_arm:
+    space: joint
+    joint_dims: [8, 9, 10, 11, 12, 13, 14]
+    joint_wrap: pi
+    gripper_dim: 15

runtime/backends/actioncodec2/assets/presets/dual_joint7_delta.yaml

@@ -0,0 +1,16 @@
+name: dual_joint7_delta
+input_dim: 16
+input: delta
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  left_arm:
+    space: joint
+    joint_dims: [0, 1, 2, 3, 4, 5, 6]
+    gripper_dim: 7
+  right_arm:
+    space: joint
+    joint_dims: [8, 9, 10, 11, 12, 13, 14]
+    gripper_dim: 15

runtime/backends/actioncodec2/assets/presets/single_eef_abs.yaml

@@ -0,0 +1,14 @@
+name: single_eef_abs
+input_dim: 7
+input: absolute
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  arm:
+    space: eef
+    position_dims: [0, 1, 2]
+    rotation_dims: [3, 4, 5]
+    rotation_format: rotation_vector
+    gripper_dim: 6

runtime/backends/actioncodec2/assets/presets/single_eef_delta.yaml

@@ -0,0 +1,14 @@
+name: single_eef_delta
+input_dim: 7
+input: delta
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  arm:
+    space: eef
+    position_dims: [0, 1, 2]
+    rotation_dims: [3, 4, 5]
+    rotation_format: rotation_vector
+    gripper_dim: 6

runtime/backends/actioncodec2/assets/presets/single_joint6_abs.yaml

@@ -0,0 +1,13 @@
+name: single_joint6_abs
+input_dim: 7
+input: absolute
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  arm:
+    space: joint
+    joint_dims: [0, 1, 2, 3, 4, 5]
+    joint_wrap: pi
+    gripper_dim: 6

runtime/backends/actioncodec2/assets/presets/single_joint6_abs_no_gripper.yaml

@@ -0,0 +1,12 @@
+name: single_joint6_abs_no_gripper
+input_dim: 6
+input: absolute
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  arm:
+    space: joint
+    joint_dims: [0, 1, 2, 3, 4, 5]
+    joint_wrap: pi

runtime/backends/actioncodec2/assets/presets/single_joint6_delta.yaml

@@ -0,0 +1,12 @@
+name: single_joint6_delta
+input_dim: 7
+input: delta
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  arm:
+    space: joint
+    joint_dims: [0, 1, 2, 3, 4, 5]
+    gripper_dim: 6

runtime/backends/actioncodec2/assets/presets/single_joint7_abs.yaml

@@ -0,0 +1,13 @@
+name: single_joint7_abs
+input_dim: 8
+input: absolute
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  arm:
+    space: joint
+    joint_dims: [0, 1, 2, 3, 4, 5, 6]
+    joint_wrap: pi
+    gripper_dim: 7

runtime/backends/actioncodec2/assets/presets/single_joint7_delta.yaml

@@ -0,0 +1,12 @@
+name: single_joint7_delta
+input_dim: 8
+input: delta
+source_hz: 15.0
+target_hz: 15.0
+interpolation: cubic
+token_order: time_major
+components:
+  arm:
+    space: joint
+    joint_dims: [0, 1, 2, 3, 4, 5, 6]
+    gripper_dim: 7

runtime/backends/actioncodec2/assets/profiles/eef.yaml

@@ -0,0 +1,74 @@
+model_type: actioncodec2_unified
+config_version: 1
+physical:
+  input_dim: 7
+  components:
+    translation_m:
+      dims:
+        start: 0
+        stop: 3
+      space: additive
+      input_order: 0
+      primitive_order: 1
+      input_transform:
+        scale: 1.0
+        clip: false
+      constraints:
+        order_0:
+          min: -2.0
+          max: 2.0
+        order_1:
+          min: -0.08
+          max: 0.08
+      distortion_limit: 0.005
+      primitive_selection: smooth
+    orientation_so3:
+      dims:
+        start: 3
+        stop: 6
+      space: so3
+      input_order: 0
+      primitive_order: 1
+      input_transform:
+        scale: 1.0
+        clip: false
+      constraints:
+        order_0:
+          min: -3.141592653589793
+          max: 3.141592653589793
+        order_1:
+          min: -0.25
+          max: 0.25
+      distortion_limit: 0.02
+      primitive_selection: smooth
+      representation: rotation_vector
+      multiplication: right
+    gripper_open_command:
+      dims:
+      - 6
+      space: additive
+      input_order: 0
+      primitive_order: 0
+      input_transform:
+        scale: 1.0
+        clip: true
+      constraints:
+        order_0:
+          min: 0.0
+          max: 1.0
+      distortion_limit: 0.5
+      primitive_selection: nearest
+set_bpe:
+  lookahead: 8
+  allow_cross_component_merge: false
+train:
+  primitives:
+    zero_primitive_components:
+    - translation_m
+    - orientation_so3
+  bpe:
+    spatial_vocab_size: 2048
+    vocab_size: 4096
+  quality_control:
+    hard_order_1_scale: 3.0
+    min_segment_length: 2

runtime/backends/actioncodec2/assets/profiles/eef_second_order.yaml

@@ -0,0 +1,80 @@
+model_type: actioncodec2_unified
+config_version: 1
+physical:
+  input_dim: 7
+  components:
+    translation_m:
+      dims:
+        start: 0
+        stop: 3
+      space: additive
+      input_order: 0
+      primitive_order: 2
+      input_transform:
+        scale: 1.0
+        clip: false
+      constraints:
+        order_0:
+          min: -2.0
+          max: 2.0
+        order_2:
+          min: -0.025
+          max: 0.025
+        order_1:
+          min: -0.08
+          max: 0.08
+      distortion_limit: 0.005
+      primitive_selection: smooth
+    orientation_so3:
+      dims:
+        start: 3
+        stop: 6
+      space: so3
+      input_order: 0
+      primitive_order: 2
+      input_transform:
+        scale: 1.0
+        clip: false
+      constraints:
+        order_0:
+          min: -3.141592653589793
+          max: 3.141592653589793
+        order_2:
+          min: -0.08
+          max: 0.08
+        order_1:
+          min: -0.25
+          max: 0.25
+      distortion_limit: 0.02
+      primitive_selection: smooth
+      representation: rotation_vector
+      multiplication: right
+    gripper_open_command:
+      dims:
+      - 6
+      space: additive
+      input_order: 0
+      primitive_order: 0
+      input_transform:
+        scale: 1.0
+        clip: true
+      constraints:
+        order_0:
+          min: 0.0
+          max: 1.0
+      distortion_limit: 0.5
+      primitive_selection: nearest
+set_bpe:
+  lookahead: 8
+  allow_cross_component_merge: false
+train:
+  primitives:
+    zero_primitive_components:
+    - translation_m
+    - orientation_so3
+  bpe:
+    spatial_vocab_size: 2048
+    vocab_size: 4096
+  quality_control:
+    hard_order_1_scale: 3.0
+    min_segment_length: 2

runtime/backends/actioncodec2/assets/profiles/joint.yaml

@@ -0,0 +1,52 @@
+model_type: actioncodec2_unified
+config_version: 1
+physical:
+  input_dim: 8
+  components:
+    generalized_joint_position:
+      dims:
+        start: 0
+        stop: 7
+      space: additive
+      input_order: 0
+      primitive_order: 1
+      input_transform:
+        scale: 1.0
+        clip: false
+      constraints:
+        order_0:
+          min: -6.283185307179586
+          max: 6.283185307179586
+        order_1:
+          min: -0.35
+          max: 0.35
+      distortion_limit: [0.004, 0.004, 0.006, 0.0075, 0.015, 0.020, 0.020]
+      primitive_selection: smooth
+    gripper_open_command:
+      dims:
+      - 7
+      space: additive
+      input_order: 0
+      primitive_order: 0
+      input_transform:
+        scale: 1.0
+        clip: true
+      constraints:
+        order_0:
+          min: 0.0
+          max: 1.0
+      distortion_limit: 0.5
+      primitive_selection: nearest
+set_bpe:
+  lookahead: 8
+  allow_cross_component_merge: false
+train:
+  primitives:
+    zero_primitive_components:
+    - generalized_joint_position
+  bpe:
+    spatial_vocab_size: 2048
+    vocab_size: 4096
+  quality_control:
+    hard_order_1_scale: 3.0
+    min_segment_length: 2

runtime/backends/actioncodec2/assets/profiles/joint_second_order.yaml

@@ -0,0 +1,62 @@
+model_type: actioncodec2_unified
+config_version: 1
+physical:
+  input_dim: 8
+  components:
+    generalized_joint_position:
+      dims:
+        start: 0
+        stop: 7
+      space: additive
+      input_order: 0
+      primitive_order: 2
+      input_transform:
+        scale: 1.0
+        clip: false
+      constraints:
+        order_0:
+          min: -6.283185307179586
+          max: 6.283185307179586
+        order_2:
+          min: -0.25
+          max: 0.25
+        order_1:
+          min: -0.35
+          max: 0.35
+      distortion_limit:
+      - 0.004
+      - 0.004
+      - 0.006
+      - 0.0075
+      - 0.015
+      - 0.02
+      - 0.02
+      primitive_selection: smooth
+    gripper_open_command:
+      dims:
+      - 7
+      space: additive
+      input_order: 0
+      primitive_order: 0
+      input_transform:
+        scale: 1.0
+        clip: true
+      constraints:
+        order_0:
+          min: 0.0
+          max: 1.0
+      distortion_limit: 0.5
+      primitive_selection: nearest
+set_bpe:
+  lookahead: 8
+  allow_cross_component_merge: false
+train:
+  primitives:
+    zero_primitive_components:
+    - generalized_joint_position
+  bpe:
+    spatial_vocab_size: 2048
+    vocab_size: 4096
+  quality_control:
+    hard_order_1_scale: 3.0
+    min_segment_length: 2

runtime/backends/actioncodec2/batch_semantics.py

@@ -0,0 +1,269 @@
+"""Canonical adapters shared by episode fitting and batched inference."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from ._so3 import (
+    conjugate_quaternion,
+    multiply_quaternions,
+    normalize_quaternion,
+    quaternion_to_rotvec,
+)
+from .native_backend import native_kernel
+from .semantic import (
+    difference_additive,
+    eef_pose_absolute,
+    eef_pose_components,
+    eef_pose_deltas,
+    eef_pose_from_components,
+    integrate_additive,
+    integrate_eef_deltas,
+    normalize_gripper_actions,
+    resample_uniform,
+    wrap_to_pi,
+)
+
+if TYPE_CHECKING:
+    from .presets import ComponentSpec, PresetConfig
+
+
+def _resample(values, source_hz, target_hz, method, *, kind="linear", target_length=None):
+    source_hz, target_hz = float(source_hz), float(target_hz)
+    if source_hz == target_hz and target_length in (None, values.shape[1]):
+        return values
+    kernel = native_kernel("resample_uniform_batch")
+    if kernel is not None:
+        return kernel(
+            values,
+            source_hz,
+            target_hz,
+            kind,
+            -1 if target_length is None else target_length,
+            method,
+        )
+    return np.stack(
+        [
+            resample_uniform(
+                row,
+                source_hz=source_hz,
+                target_hz=target_hz,
+                kind=kind,
+                target_length=target_length,
+                method=method,
+            )
+            for row in values
+        ]
+    )
+
+
+def _pose_map(values, kernel_name, reference):
+    kernel = native_kernel(kernel_name)
+    return kernel(values) if kernel is not None else np.stack([reference(row) for row in values])
+
+
+def _rotvec_poses(values):
+    kernel = native_kernel("eef_pose_from_rotvec_batch")
+    if kernel is not None:
+        return kernel(values)
+    flat = values.reshape(-1, 6)
+    return eef_pose_from_components(flat[:, :3], flat[:, 3:], "rotation_vector").reshape(
+        *values.shape[:2], 7
+    )
+
+
+def _poses(values: np.ndarray, component: ComponentSpec) -> np.ndarray:
+    position = np.take(values, component.position_dims, axis=-1)
+    rotation = np.take(values, component.rotation_dims, axis=-1)
+    if component.rotation_format == "rotation_vector":
+        return _rotvec_poses(np.concatenate((position, rotation), axis=-1))
+    return eef_pose_from_components(
+        position.reshape(-1, 3), rotation.reshape(-1, rotation.shape[-1]), component.rotation_format
+    ).reshape(*values.shape[:2], 7)
+
+
+def _local(poses, anchor):
+    quaternion = normalize_quaternion(poses[..., 3:7])
+    inverse = conjugate_quaternion(normalize_quaternion(anchor[..., 3:7]))
+    relative = multiply_quaternions(inverse, quaternion)
+    return np.concatenate(
+        (poses[..., :3].astype(np.float64) - anchor[..., :3], quaternion_to_rotvec(relative)),
+        axis=-1,
+    ).astype(np.float32)
+
+
+def _joint(preset, values, component, source_hz, current_state):
+    source = np.take(values, component.joint_dims, axis=-1)
+    if preset.input == "delta":
+        source = integrate_additive(source, 1)
+        anchor = np.zeros((len(source), 1, source.shape[-1]), np.float32)
+    else:
+        anchor = (
+            source[:, :1]
+            if current_state is None
+            else np.take(current_state, component.joint_dims, axis=-1)[:, None]
+        )
+        joined = np.concatenate((anchor, source), axis=1)
+        if component.joint_wrap == "pi":
+            outside = np.abs(joined) > np.pi + 1e-5
+            if outside.any():
+                row = int(np.flatnonzero(outside.any(axis=(1, 2)))[0])
+                column = component.joint_dims[int(np.argmax(outside[row].any(axis=0)))]
+                raise ValueError(
+                    f"action space {preset.name!r} declares joint_wrap='pi' for component "
+                    f"{component.name!r}, so its absolute joint columns must lie in (-pi, pi]; "
+                    f"source column {column} reaches {joined[row][outside[row]][0]:.6g} rad. Wrap the "
+                    "inputs, or register a layout with joint_wrap='none' for coordinates that "
+                    "are continuous past one turn, such as a multi-turn or prismatic axis."
+                )
+        batch, steps, width = joined.shape
+        kernel = native_kernel("unwrap_angles")
+        unwrapped = (
+            (
+                kernel(joined.transpose(1, 0, 2).reshape(steps, batch * width), True)
+                .reshape(steps, batch, width)
+                .transpose(1, 0, 2)
+            )
+            if kernel is not None
+            else np.unwrap(joined, axis=1)
+        )
+        anchor, source = unwrapped[:, :1], unwrapped[:, 1:]
+    return np.asarray(
+        _resample(source, source_hz, preset.target_hz, preset.interpolation) - anchor,
+        dtype=np.float32,
+    )
+
+
+def _eef(preset, values, component, source_hz, current_state):
+    if preset.input == "delta":
+        delta = np.concatenate(
+            (
+                np.take(values, component.position_dims, axis=-1),
+                np.take(values, component.rotation_dims, axis=-1),
+            ),
+            axis=-1,
+        )
+        poses = _pose_map(delta, "integrate_eef_deltas_batch", integrate_eef_deltas)
+        sampled = _resample(
+            poses, source_hz, preset.target_hz, preset.interpolation, kind="pose_xyzw"
+        )
+        return _pose_map(sampled, "eef_pose_absolute_batch", eef_pose_absolute)
+    poses = _poses(values, component)
+    anchor = poses[:, :1] if current_state is None else _poses(current_state[:, None], component)
+    sampled = _resample(poses, source_hz, preset.target_hz, preset.interpolation, kind="pose_xyzw")
+    return _local(sampled, anchor)
+
+
+def canonicalize_batch(preset: PresetConfig, values, *, source_hz, current_state):
+    """Return component ``(B,T_codec,D_active)`` arrays from validated input."""
+    return {
+        component.name: canonicalize_component_batch(
+            preset, values, component, source_hz=source_hz, current_state=current_state
+        )
+        for component in preset.components
+    }
+
+
+def canonicalize_component_batch(preset, values, component, *, source_hz, current_state):
+    """Convert one component of validated ``(B,T,D)`` source actions."""
+    motion = (_joint if component.space == "joint" else _eef)(
+        preset, values, component, source_hz, current_state
+    )
+    if component.gripper_dim is None:
+        return motion
+    gripper = values[..., component.gripper_dim : component.gripper_dim + 1]
+    if component.gripper_threshold is not None:
+        gripper = ((gripper >= component.gripper_threshold) if component.gripper_open_is_high
+                   else (gripper <= component.gripper_threshold)).astype(np.float32)
+    else:
+        gripper = normalize_gripper_actions(
+            gripper.reshape(-1, 1),
+            closed_value=component.gripper_closed_value,
+            open_value=component.gripper_open_value,
+        ).reshape(*values.shape[:2], 1)
+    gripper = _resample(gripper, source_hz, preset.target_hz, preset.interpolation, kind="hold")
+    return np.concatenate((motion, gripper), axis=-1)
+
+
+def _external_eef(motion, component, current_state):
+    poses = _rotvec_poses(motion)
+    if current_state is None:
+        return _pose_map(poses, "eef_pose_deltas_batch", eef_pose_deltas)
+    anchor = _poses(current_state[:, None], component)
+    quaternion = multiply_quaternions(
+        normalize_quaternion(anchor[..., 3:7]), normalize_quaternion(poses[..., 3:7])
+    )
+    attached = np.concatenate(
+        (poses[..., :3].astype(np.float64) + anchor[..., :3], quaternion), axis=-1
+    ).astype(np.float32)
+    return eef_pose_components(attached.reshape(-1, 7), component.rotation_format).reshape(
+        *attached.shape[:2], 3 + len(component.rotation_dims)
+    )
+
+
+def externalize_batch(
+    values, component: ComponentSpec, *, source_hz, target_hz, target_length, method, current_state
+):
+    """Resample decoded paths and return ``(B,T_output,D_component)`` arrays."""
+    sampled = resample_canonical_batch(
+        values,
+        component,
+        source_hz=source_hz,
+        target_hz=target_hz,
+        target_length=target_length,
+        method=method,
+    )
+    return externalize_canonical_batch(sampled, component, current_state=current_state)
+
+
+def resample_canonical_batch(
+    values, component: ComponentSpec, *, source_hz, target_hz, target_length, method
+):
+    """Rate-convert local ``(B,T,D_active)`` paths, holding gripper commands."""
+    motion_width = len(component.joint_dims or ()) if component.space == "joint" else 6
+    motion = values[..., :motion_width]
+    if component.space == "eef":
+        poses = _rotvec_poses(motion)
+        poses = _resample(
+            poses, source_hz, target_hz, method, kind="pose_xyzw", target_length=target_length
+        )
+        motion = _pose_map(poses, "eef_pose_absolute_batch", eef_pose_absolute)
+    else:
+        motion = _resample(motion, source_hz, target_hz, method, target_length=target_length)
+    if component.gripper_dim is None:
+        return np.asarray(motion, dtype=np.float32)
+    gripper = _resample(
+        values[..., motion_width : motion_width + 1],
+        source_hz,
+        target_hz,
+        method,
+        kind="hold",
+        target_length=target_length,
+    )
+    return np.asarray(np.concatenate((motion, gripper), axis=-1), dtype=np.float32)
+
+
+def externalize_canonical_batch(values, component: ComponentSpec, *, current_state):
+    """Restore source coordinates or deltas from ``(B,T,D_active)`` local paths."""
+    motion_width = len(component.joint_dims or ()) if component.space == "joint" else 6
+    motion = values[..., :motion_width]
+    if component.space == "eef":
+        motion = _external_eef(motion, component, current_state)
+    elif current_state is None:
+        motion = difference_additive(motion, 1)
+    else:
+        motion = motion + np.take(current_state, component.joint_dims, axis=-1)[:, None]
+        if component.joint_wrap == "pi":
+            motion = wrap_to_pi(motion)
+    if component.gripper_dim is None:
+        return np.asarray(motion, dtype=np.float32)
+    gripper = values[..., motion_width : motion_width + 1]
+    if component.gripper_threshold is not None:
+        gripper = gripper if component.gripper_open_is_high else 1.0 - gripper
+        return np.asarray(np.concatenate((motion, gripper), axis=-1), dtype=np.float32)
+    gripper = component.gripper_closed_value + gripper * (
+        component.gripper_open_value - component.gripper_closed_value
+    )
+    return np.asarray(np.concatenate((motion, gripper), axis=-1), dtype=np.float32)