initium-deeplearning 0.1.0__py3-none-any.whl

initium/data/multimodal/multimodal_dataset.py

@@ -0,0 +1,401 @@
+"""
+file: data/multimodal/multimodal_dataset.py
+
+Fuses two or more of {text, audio, video} chain datasets. The SAME (key,
+value) facts/queries draw is shown through EACH active modality's own
+encoding channel simultaneously, concatenated into one wider per-timestep
+input vector (input_dim = sum of active modalities' input_dims). This is
+what makes "which modality does the model actually use" answerable via a
+direct ablation: zero one modality's channel slice and re-evaluate -
+Delta_<modality>, the Dependency-layer (Layer C) analogue of
+evaluate_id_ablated (memory) / evaluate_id_combiner_stage_ablated
+(component), for modalities.
+
+Only the FIRST listed modality's targets are scored ("primary") - this
+tests whether auxiliary modalities help/are relied upon for the primary
+modality's own recall+chain task, not multi-task multi-output prediction.
+
+Real-data wiring: each modality in the `modalities` list passed to
+__init__ may be a bare name (synthetic, unchanged default) or
+"modality:path" to source that modality's own real fact pool from `path`
+via its existing single-modality real-data pipeline (data/common/
+real_data.py) -- see _parse_modality_specs() below. Repeating the SAME
+path across modalities (e.g. "video:clip.mp4+audio:clip.mp4") draws both
+modalities' real content from one shared file, such as a video that also
+carries an audio track -- the flagship "upload one video, learn its
+picture AND its soundtrack" case. See MultimodalDataset.__init__'s
+`link_role` for how training vs. inference (MultimodalTask) route a
+spec's path to the right underlying *ChainDataset kwarg.
+"""
+
+import random
+
+import torch
+
+from initium.data.audio.audio_dataset import AudioChainDataset
+from initium.data.base_dataset import BaseDataset
+from initium.data.common.chain_task import (
+    ChainCurriculum,
+    score_answer_steps,
+    value_cumsum_field_log_header,
+    write_value_cumsum_field_log,
+)
+from initium.data.text.text_dataset import TextChainDataset
+from initium.data.video.video_dataset import VideoChainDataset
+
+_MODALITY_CLASSES = {
+    "text": TextChainDataset,
+    "audio": AudioChainDataset,
+    "video": VideoChainDataset,
+}
+
+
+def _parse_modality_specs(modalities):
+    """Each entry of `modalities` is either a bare modality name ("text",
+    "audio", "video") -- synthetic facts, unchanged default -- or
+    "modality:path" to source that modality's OWN real (key,value) fact
+    pool from `path`, via the exact same per-modality real-data pipeline
+    single-modality datasets already use (data/common/real_data.py). The
+    SAME path can be given to more than one modality -- e.g.
+    "video:clip.mp4+audio:clip.mp4" -- to draw both modalities' real
+    content from one shared file, such as a video that also carries an
+    audio track: video_token_stream (opencv) and audio_token_stream
+    (torchaudio) each open the container with their own library and only
+    ever read the stream they care about, so pointing both at the same
+    path is safe and is the intended way to say "learn from this video's
+    picture AND its soundtrack." -> [(name, path_or_None), ...].
+    """
+    parsed = []
+    for spec in modalities:
+        spec = (spec or "").strip()
+        if not spec:
+            continue
+        name, _, path = spec.partition(":")
+        parsed.append((name.strip().lower(), path.strip() or None))
+    return parsed
+
+
+class MultimodalDataset(BaseDataset):
+    # Real-data wiring: see _parse_modality_specs() above for the
+    # "modality[:path]" spec syntax and MultimodalTask (inference/tasks/
+    # multimodal_task.py) for how --dataset-link reaches here at inference.
+    def __init__(self, modalities, link_role: str = "dataset_link"):
+        """
+        modalities: iterable of modality specs (see _parse_modality_specs).
+        link_role: "dataset_link" (default; training-side) -- a spec's path
+            becomes that modality's OWN training source (auto-split into a
+            disjoint train/test pool by that modality's dataset class,
+            exactly like a standalone *ChainDataset(dataset_link=...)), or
+            "test_dataset_link" (inference-side, see MultimodalTask) -- a
+            spec's path becomes a fixed, already-real TEST pool directly.
+        """
+        if link_role not in ("dataset_link", "test_dataset_link"):
+            raise ValueError(
+                "MultimodalDataset: link_role must be 'dataset_link' "
+                f"or 'test_dataset_link', got {link_role!r}"
+            )
+        specs = _parse_modality_specs(modalities)
+        if len(specs) < 2:
+            raise ValueError(f"MultimodalDataset needs >=2 modalities, got {modalities!r}")
+        names = [n for n, _ in specs]
+        unknown = [n for n in names if n not in _MODALITY_CLASSES]
+        if unknown:
+            raise ValueError(
+                f"MultimodalDataset: unknown modalities {unknown}, "
+                f"expected a subset of {sorted(_MODALITY_CLASSES)}"
+            )
+        self.modalities = names
+        self.subs = [
+            _MODALITY_CLASSES[name](**({link_role: path} if path else {})) for name, path in specs
+        ]
+        self.primary = self.subs[0]
+        self.name = "multimodal[" + "+".join(names) + "]"
+
+        off = 0
+        self._offsets = []
+        for s in self.subs:
+            self._offsets.append((off, off + s.input_dim))
+            off += s.input_dim
+        self.input_dim = off
+        self.output_dim = self.primary.output_dim
+        self.eval_batch_size = self.primary.eval_batch_size
+        self.advance_threshold = self.primary.advance_threshold
+        self._table = self.primary._table
+        self._lesson_nr_cells = self.primary._lesson_nr_cells
+
+    def _facts_pool(self, use_test_pool: bool = False):
+        """Real (key,value) pool to draw a fused episode from, or None for
+        the fully-synthetic draw (unchanged default when no modality was
+        given a real source). use_test_pool selects each sub's held-out
+        `_test_fact_pool` (evaluate_ood) vs its training `_fact_pool`
+        (sample_batch / evaluate_id_ablated / evaluate_modality_ablated) --
+        same train/test convention every single-modality dataset already
+        uses. Values always come from the PRIMARY modality's pool (only the
+        primary is scored -- see class docstring); keys are restricted to
+        ones ALSO present in every OTHER real-sourced modality's pool, so
+        any channel built from real data genuinely has real content for
+        every key an episode can draw (e.g. a video file's frame-derived
+        pool intersected with that SAME file's audio-track-derived pool).
+        Falls back to the unrestricted primary pool if the intersection is
+        empty, rather than raising mid-run."""
+        attr = "_test_fact_pool" if use_test_pool else "_fact_pool"
+        primary_pool = getattr(self.primary, attr)
+        if primary_pool is None:
+            return None
+        other_pools = [p for p in (getattr(s, attr) for s in self.subs[1:]) if p is not None]
+        if not other_pools:
+            return primary_pool
+        shared_keys = set(k for k, _ in primary_pool)
+        for p in other_pools:
+            shared_keys &= set(k for k, _ in p)
+        restricted = [(k, v) for k, v in primary_pool if k in shared_keys]
+        return restricted or primary_pool
+
+    def make_curriculum(self):
+        return ChainCurriculum(self, self._table, self._lesson_nr_cells)
+
+    def set_output_proj(self, proj):
+        self._output_proj = proj
+
+    def _build_fused_episode(
+        self,
+        num_facts_range,
+        num_queries_range,
+        rng=None,
+        skip_modalities=None,
+        perturb=None,
+        use_test_pool=False,
+    ):
+        rng = rng if rng is not None else random
+        skip_modalities = skip_modalities or set()
+        codec, label_range = self.primary.codec, self.primary.label_range
+        NF = codec.num_digits
+
+        def _as_range(x):
+            return x if isinstance(x, tuple) else (x, x)
+
+        num_facts_range = _as_range(num_facts_range)
+        num_queries_range = _as_range(num_queries_range)
+
+        num_facts = rng.randint(*num_facts_range)
+        pool = self._facts_pool(use_test_pool=use_test_pool)
+        if pool is not None:
+            num_facts = min(num_facts, len(pool))
+            facts = rng.sample(pool, num_facts)
+        else:
+            keys = rng.sample(range(label_range), num_facts)
+            vals = [rng.randrange(label_range) for _ in range(num_facts)]
+            facts = list(zip(keys, vals))
+        rng.shuffle(facts)
+        keys = [k for k, _ in facts]
+        num_queries = rng.randint(*num_queries_range)
+        query_keys = [rng.choice(keys) for _ in range(num_queries)]
+
+        def fused(build_fns):
+            return torch.cat(
+                [
+                    torch.zeros(sub.input_dim) if i in skip_modalities else build_fns[i]()
+                    for i, sub in enumerate(self.subs)
+                ]
+            )
+
+        inputs, target_digits, answer_mask = [], [], []
+        for i, (k, v) in enumerate(facts):
+            pt = 1.0 if i == 0 else 0.0
+            inputs.append(
+                fused(
+                    [
+                        (lambda s=sub: s._encode_fact(k, v, pt, perturb=perturb, rng=rng))
+                        for sub in self.subs
+                    ]
+                )
+            )
+            target_digits.append([0] * (2 * NF))
+            answer_mask.append(0)
+        for i, k in enumerate(query_keys):
+            pt = 1.0 if i == 0 else 0.0
+            inputs.append(fused([(lambda s=sub: s._encode_query(k, pt)) for sub in self.subs]))
+            target_digits.append([0] * (2 * NF))
+            answer_mask.append(0)
+        value_of, cumsum = dict(facts), 0
+        for i, k in enumerate(query_keys):
+            cumsum = (cumsum + value_of[k]) % label_range
+            pt = 1.0 if i == 0 else 0.0
+            inputs.append(fused([(lambda s=sub: s._encode_answer(pt)) for sub in self.subs]))
+            target_digits.append(
+                codec.label_to_digit_targets(value_of[k]) + codec.label_to_digit_targets(cumsum)
+            )
+            answer_mask.append(1)
+        return (
+            torch.stack(inputs),
+            torch.tensor(target_digits, dtype=torch.long),
+            torch.tensor(answer_mask, dtype=torch.float32),
+            num_queries,
+        )
+
+    def collate(self, batch):
+        max_len = max(len(x[0]) for x in batch)
+        B = len(batch)
+        NF = self.primary.codec.num_digits
+        padded_input = torch.zeros(B, max_len, self.input_dim)
+        padded_targets = torch.zeros(B, max_len, 2 * NF, dtype=torch.long)
+        padded_mask = torch.zeros(B, max_len)
+        for i, (inp, tgt, mask, _) in enumerate(batch):
+            T = inp.size(0)
+            padded_input[i, :T] = inp
+            padded_targets[i, :T] = tgt
+            padded_mask[i, :T] = mask
+        return padded_input, padded_targets, padded_mask
+
+    def sample_batch(self, curriculum, batch_size):
+        episodes = []
+        for _ in range(batch_size):
+            idx = (
+                random.randint(0, curriculum.lesson - 1)
+                if curriculum.lesson > 0 and random.random() < 0.10
+                else curriculum.lesson
+            )
+            nf, nq = curriculum.table[idx]
+            episodes.append(self._build_fused_episode(nf, nq))
+        return self.collate(episodes)  # was self.primary.collate(episodes)
+
+    def loss(self, output, target, mask):
+        return self.primary.loss(output, target, mask)
+
+    def diversity(self, output, target, mask):
+        return self.primary.diversity(output, target, mask)
+
+    def evaluate_chain(
+        self,
+        model,
+        device,
+        num_episodes,
+        num_facts_range=None,
+        num_queries_range=None,
+        rng=None,
+        ablate_memory=False,
+        step_breakdown=False,
+        skip_modalities=None,
+        verbose_n=0,
+        use_test_pool=False,
+        skip_stages=None,
+        field_log=None,
+    ):
+        model.eval()
+        total = correct_value = correct_cumsum = perfect_episodes = tested = 0
+        by_depth = {}
+        codec = self.primary.codec
+        with torch.no_grad():
+            while tested < num_episodes:
+                input_seq, target_digits, answer_mask, depth = self._build_fused_episode(
+                    num_facts_range,
+                    num_queries_range,
+                    rng=rng,
+                    skip_modalities=skip_modalities,
+                    use_test_pool=use_test_pool,
+                )
+                input_seq = input_seq.unsqueeze(0).to(device)
+                output, _ = model(
+                    input_seq,
+                    (None, None, None),
+                    reset_experience=True,
+                    pass_through_memory=not ablate_memory,
+                    **({"combiner_skip_stages": skip_stages} if skip_stages else {}),
+                )
+                output = self._output_proj(output.transpose(0, 1).contiguous().squeeze(0))
+                ep_total, ep_correct, episode_perfect, ep_cv, ep_cc, ep_n = score_answer_steps(
+                    output,
+                    target_digits,
+                    answer_mask,
+                    codec,
+                    verbose=(tested < verbose_n),
+                    field_log=field_log,
+                )
+                correct_value += ep_cv
+                correct_cumsum += ep_cc
+                total += ep_n
+                perfect_episodes += int(episode_perfect)
+                tested += 1
+                if step_breakdown:
+                    acc = by_depth.setdefault(ep_total, [0, 0, 0, 0])
+                    acc[0] += 2 * ep_total
+                    acc[1] += ep_correct
+                    acc[2] += 1
+                    acc[3] += int(episode_perfect)
+        combined_acc = 100.0 * (correct_value + correct_cumsum) / max(2 * total, 1)
+        perfect_frac = 100.0 * perfect_episodes / max(tested, 1)
+        print(
+            f"Eval [{self.name}]: value {100.0 * correct_value / max(total, 1):.2f}% | "
+            f"cumsum {100.0 * correct_cumsum / max(total, 1):.2f}% | combined {combined_acc:.2f}% | "
+            f"perfect {perfect_frac:.2f}% ({tested} episodes)"
+        )
+        model.train()
+        if step_breakdown:
+            breakdown = {
+                d: (100.0 * c / max(t, 1), 100.0 * p / max(e, 1), e)
+                for d, (t, c, e, p) in sorted(by_depth.items())
+            }
+            return combined_acc, perfect_frac, breakdown
+        return combined_acc, perfect_frac
+
+    def evaluate_ood(self, model, device, num_episodes, rng, verbose_n=0, field_log=None):
+        return self.evaluate_chain(
+            model,
+            device,
+            num_episodes,
+            num_facts_range=(15, 20),
+            num_queries_range=self.primary.ood_query_range,
+            rng=rng,
+            step_breakdown=True,
+            use_test_pool=True,
+            verbose_n=verbose_n,
+            field_log=field_log,
+        )
+
+    def evaluate_id_ablated(self, model, device, curriculum, lesson_idx):
+        nf, nq = curriculum.table[lesson_idx]
+        return self.evaluate_chain(
+            model,
+            device,
+            self.eval_batch_size,
+            num_facts_range=nf,
+            num_queries_range=nq,
+            ablate_memory=True,
+        )
+
+    def evaluate_id_combiner_stage_ablated(
+        self, model, device, curriculum, lesson_idx, skip_stages
+    ):
+        """Combiner-stage-off eval on lesson `lesson_idx` -> (acc_pct, perfect_pct).
+        Mirrors data/common/chain_task.py's KVChainDataset implementation --
+        only meaningful when model.combiner_wrapper is a ChainedControllerWrapper
+        (e.g. --split-graph-combiner-variant mamba+cfc)."""
+        nf, nq = curriculum.table[lesson_idx]
+        return self.evaluate_chain(
+            model,
+            device,
+            self.eval_batch_size,
+            num_facts_range=nf,
+            num_queries_range=nq,
+            skip_stages=skip_stages,
+        )
+
+    def evaluate_modality_ablated(self, model, device, curriculum, lesson_idx, modality: str):
+        """Delta_<modality> = full_acc - ablated_acc, per the research doc's
+        modality-dependency proposal."""
+        idx = self.modalities.index(modality)
+        nf, nq = curriculum.table[lesson_idx]
+        return self.evaluate_chain(
+            model,
+            device,
+            self.eval_batch_size,
+            num_facts_range=nf,
+            num_queries_range=nq,
+            skip_modalities={idx},
+        )
+
+    def field_log_header(self):
+        return value_cumsum_field_log_header()
+
+    def write_field_log(self, writer, file, step, lesson, eval_type, field_log):
+        write_value_cumsum_field_log(writer, file, step, lesson, eval_type, field_log)

initium/data/text/text_dataset.py

@@ -0,0 +1,79 @@
+"""
+file: data/text/text_dataset.py
+
+Text modality dataset: TextChainDataset(KVChainDataset). Default (no
+dataset_link): "facts" are token(key)->token(value) associations drawn
+fresh each episode from the full label space (unchanged synthetic
+behavior). With dataset_link (a path to a real text file): a byte-level
+BPE tokenizer (data/common/real_data.py) is trained on the file, the
+resulting token stream is turned into a real (key,value) fact pool, and
+that pool is split into a disjoint TRAIN pool (used by build_episode
+instead of random draws) and TEST pool (used by build_ood_facts instead of
+the synthetic seeded table below) -- the text analogue of the graph
+dataset's "train on synthetic graphs, test on the held-out London
+Underground" split. test_dataset_link (a second, separate text file) can
+be given to use an entirely different source as the test set instead of a
+held-out slice of the same file.
+
+Depth axis: number of queries (reasoning-chain-length / context-length
+analogue). Fields: value (lookup), cumsum (chained state).
+"""
+
+import random
+
+from initium.data.common.chain_task import KVChainDataset
+from initium.data.common.real_data import build_text_kv_pool, split_train_test_facts
+
+TEXT_CURRICULUM = [
+    (3, 2),
+    (3, 3),
+    (5, 3),
+    (5, 4),
+    (8, 4),
+    (8, 5),
+    (10, 5),
+    (10, 6),
+    (12, 6),
+    (12, 7),
+    (15, 8),
+    (15, 9),
+    (20, 10),
+    (20, 12),
+]
+TEXT_LESSON_NR_CELLS = [128, 128, 128, 128, 160, 160, 160, 160, 192, 192, 192, 192, 256, 256]
+assert len(TEXT_CURRICULUM) == len(TEXT_LESSON_NR_CELLS)
+
+
+class TextChainDataset(KVChainDataset):
+    name = "text-chain"
+
+    def __init__(self, dataset_link: str | None = None, test_dataset_link: str | None = None):
+        self._table = TEXT_CURRICULUM
+        self._lesson_nr_cells = TEXT_LESSON_NR_CELLS
+        self._fact_pool = None
+        self._test_fact_pool = None
+        if dataset_link is not None:
+            pool = build_text_kv_pool(dataset_link, self.label_range)
+            if test_dataset_link is not None:
+                self._fact_pool = pool
+                self._test_fact_pool = build_text_kv_pool(test_dataset_link, self.label_range)
+            else:
+                self._fact_pool, self._test_fact_pool = split_train_test_facts(pool)
+            print(
+                f"[text-chain] {len(self._fact_pool)} train facts from {dataset_link} | "
+                f"{len(self._test_fact_pool)} test facts"
+                + (f" from {test_dataset_link}" if test_dataset_link else " (held-out split)")
+            )
+        elif test_dataset_link is not None:
+            # Inference-only: no training pool needed, just a fixed real test
+            # pool -- every inference/tasks/*_task.py passes --dataset-link
+            # through as test_dataset_link, since inference only ever tests
+            # an already-trained model (see this class's own docstring).
+            self._test_fact_pool = build_text_kv_pool(test_dataset_link, self.label_range)
+            print(f"[text-chain] {len(self._test_fact_pool)} test facts from {test_dataset_link}")
+
+    def _synthetic_ood_facts(self, n_facts: int = 20):
+        rng = random.Random(4242)  # fixed held-out vocabulary
+        keys = rng.sample(range(self.label_range), n_facts)
+        vals = [rng.randrange(self.label_range) for _ in range(n_facts)]
+        return list(zip(keys, vals))

initium/data/video/video_dataset.py

@@ -0,0 +1,85 @@
+"""
+file: data/video/video_dataset.py
+
+Video modality dataset: VideoChainDataset(KVChainDataset). Default (no
+dataset_link): synthetic object-identity(key)->action/state-identity(value)
+facts, drawn fresh each episode (unchanged behavior). With dataset_link (a
+path to a real video file, needs opencv-python): each frame is downsampled
+and quantized to an intensity-bin id (data/common/real_data.py
+video_token_stream) and turned into a real fact pool, split into a
+disjoint TRAIN/TEST pool the same way text_dataset.py does -- the video
+analogue of the graph dataset's held-out London Underground graph.
+test_dataset_link (a second video file) can be given to use a genuinely
+different clip as the test set instead of a held-out slice of the same
+one.
+
+Depth axis: number of queried events (event-chain-length / temporal-
+horizon analogue). Fields: value (object-action binding accuracy), cumsum
+(chained event-ordering state).
+Robustness axis: frame dropping / occlusion -- the WHOLE fact observation
+(key AND value) is zeroed with prob=severity, rather than corrupted in
+place (closer to "this frame was dropped/occluded" than to token noise).
+"""
+
+import random
+
+import torch
+
+from initium.data.common.chain_task import KVChainDataset
+from initium.data.common.real_data import build_video_kv_pool, split_train_test_facts
+
+VIDEO_CURRICULUM = [
+    (3, 2),
+    (3, 3),
+    (5, 3),
+    (5, 4),
+    (8, 4),
+    (8, 5),
+    (10, 5),
+    (10, 6),
+    (12, 6),
+    (12, 7),
+    (15, 8),
+    (15, 9),
+    (20, 10),
+    (20, 12),
+]
+VIDEO_LESSON_NR_CELLS = [128, 128, 128, 128, 160, 160, 160, 160, 192, 192, 192, 192, 256, 256]
+assert len(VIDEO_CURRICULUM) == len(VIDEO_LESSON_NR_CELLS)
+
+
+class VideoChainDataset(KVChainDataset):
+    name = "video-chain"
+
+    def __init__(self, dataset_link: str | None = None, test_dataset_link: str | None = None):
+        self._table = VIDEO_CURRICULUM
+        self._lesson_nr_cells = VIDEO_LESSON_NR_CELLS
+        self._fact_pool = None
+        self._test_fact_pool = None
+        if dataset_link is not None:
+            pool = build_video_kv_pool(dataset_link, self.label_range)
+            if test_dataset_link is not None:
+                self._fact_pool = pool
+                self._test_fact_pool = build_video_kv_pool(test_dataset_link, self.label_range)
+            else:
+                self._fact_pool, self._test_fact_pool = split_train_test_facts(pool)
+            print(
+                f"[video-chain] {len(self._fact_pool)} train facts from {dataset_link} | "
+                f"{len(self._test_fact_pool)} test facts"
+                + (f" from {test_dataset_link}" if test_dataset_link else " (held-out split)")
+            )
+        elif test_dataset_link is not None:
+            # Inference-only real test pool -- see text_dataset.py's twin fix.
+            self._test_fact_pool = build_video_kv_pool(test_dataset_link, self.label_range)
+            print(f"[video-chain] {len(self._test_fact_pool)} test facts from {test_dataset_link}")
+
+    def _synthetic_ood_facts(self, n_facts: int = 20):
+        rng = random.Random(31337)  # unseen object/action table
+        keys = rng.sample(range(self.label_range), n_facts)
+        vals = [rng.randrange(self.label_range) for _ in range(n_facts)]
+        return list(zip(keys, vals))
+
+    def perturb_fact(self, key_vec: torch.Tensor, value_vec: torch.Tensor, rng, severity: float):
+        if rng.random() < severity:
+            return torch.zeros_like(key_vec), torch.zeros_like(value_vec)
+        return key_vec, value_vec

initium/inference/__init__.py

@@ -0,0 +1 @@
+"""Inference package: reusable, checkpoint-driven inference for the DNC project."""

initium/inference/cache/__init__.py

@@ -0,0 +1 @@
+"""Inference caches (prefix state / result), modular like data/ and inference/tasks/."""

initium/inference/cache/base_cache.py

@@ -0,0 +1,112 @@
+"""
+file: inference/cache/base_cache.py
+
+Interface the cached engine talks to (counterpart of data/base_dataset.py). A new
+cache type subclasses BaseCache, overrides only the hooks it needs, and registers
+itself in cache_registry.py; the engine never imports anything cache-specific.
+
+Two hook families:
+  episode-level : lookup_episode() / store_episode()   (wrap the whole episode)
+  prefix-level  : find_resume() / store_boundary()     (accelerate the computation;
+                  only used when the caches' `wants_boundaries` is True and the
+                  episode declares cache_boundaries)
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from initium.inference.cache.keys import digest, hash_tensor
+from initium.inference.cache.state_utils import tree_nbytes
+from initium.inference.cache.stores import CacheStats, TieredStore
+
+
+@dataclass
+class EpisodeCtx:
+    index: int
+    episode: Any  # inference.tasks.base_task.Episode
+    reset_experience: bool
+    in_hash: str  # hash of the episode's full input sequence
+    chain: str  # "root" in fresh mode; history hash in persistent mode
+    boundaries: list[int] = field(default_factory=list)  # validated cut points
+    _prefix_hashes: dict[int, str] = field(default_factory=dict, repr=False)
+
+    def prefix_hash(self, boundary: int) -> str:
+        h = self._prefix_hashes.get(boundary)
+        if h is None:
+            h = self._prefix_hashes[boundary] = hash_tensor(self.episode.input_seq[:boundary])
+        return h
+
+
+@dataclass
+class EpisodeHit:
+    output: torch.Tensor  # (T, output_dim) CPU float
+    hidden: Any  # restored state (persistent mode) or None
+    saved_ms: float
+
+
+@dataclass
+class ResumePoint:
+    position: int  # steps [0, position) are already reflected in `hidden`
+    hidden: Any
+    saved_ms: float
+    cache_name: str
+
+
+class BaseCache:
+    name: str = "base"
+    wants_boundaries: bool = False
+
+    def __init__(self, store: TieredStore, model_fp: str, device: torch.device):
+        self.store, self.model_fp, self.device = store, model_fp, device
+        self.stats = CacheStats()
+
+    # ---- hooks ----------------------------------------------------------------
+    def applicable(self, *, reset_experience: bool, resumable: bool) -> tuple[bool, str]:
+        return True, ""
+
+    def lookup_episode(self, ctx: EpisodeCtx) -> EpisodeHit | None:
+        return None
+
+    def store_episode(
+        self, ctx: EpisodeCtx, output: torch.Tensor, hidden_after: Any, compute_ms: float
+    ) -> None:
+        pass
+
+    def find_resume(self, ctx: EpisodeCtx) -> ResumePoint | None:
+        return None
+
+    def store_boundary(
+        self, ctx: EpisodeCtx, boundary: int, hidden: Any, compute_ms: float
+    ) -> None:
+        pass
+
+    # ---- helpers --------------------------------------------------------------
+    def report(self, reset: bool = False) -> dict:
+        d = self.stats.to_dict()
+        if reset:
+            self.stats = CacheStats()
+        return d
+
+    def _key(self, *parts: Any) -> str:
+        return f"{self.name}-{digest(self.model_fp, self.name, *parts)}"
+
+    def _get(self, key: str):
+        return self.store.get(key)
+
+    def _put(self, key: str, payload: Any, compute_ms: float) -> None:
+        nbytes = tree_nbytes(payload)
+        self.store.put(
+            key,
+            {
+                "payload": payload,
+                "compute_ms": float(compute_ms),
+                "nbytes": nbytes,
+                "kind": self.name,
+            },
+        )
+        self.stats.stores += 1
+        self.stats.bytes_stored += nbytes