source.py

"""The literal single trainable source behind every Dendro operation."""

from __future__ import annotations

import hashlib
import math
from dataclasses import dataclass
from functools import lru_cache
from typing import Any, Iterable, Mapping, Sequence

import torch
from torch import nn
from torch.nn import functional as F


@dataclass(frozen=True, slots=True)
class PackedTernaryTensor:
    """Two-bit storage for an exported 1.58-bit tensor.

    Codes are packed four per byte: 0 -> zero, 1 -> positive, 2 -> negative.
    The trainable model remains a single floating-point source; this object is an
    inference/export representation and is intentionally not an ``nn.Parameter``.
    """

    packed: torch.Tensor
    scale: torch.Tensor
    shape: tuple[int, ...]
    numel: int

    def unpack(self, *, device: torch.device | str | None = None, dtype: torch.dtype | None = None) -> torch.Tensor:
        raw = self.packed.to(device=device)
        shifts = torch.tensor([0, 2, 4, 6], device=raw.device, dtype=torch.uint8)
        codes = ((raw.unsqueeze(-1) >> shifts) & 0b11).reshape(-1)[: self.numel].reshape(self.shape)
        scale = self.scale.to(device=raw.device)
        out = torch.zeros(self.shape, device=raw.device, dtype=scale.dtype)
        out = torch.where(codes == 1, scale, out)
        out = torch.where(codes == 2, -scale, out)
        return out.to(dtype=dtype or self.scale.dtype)


class DendroSourceLayer(nn.Module):
    """One registered trainable tensor serving the complete architecture.

    Every logical weight is a deterministic differentiable gather from ``source``.
    Multiple logical systems therefore share and co-train one substrate instead of
    owning separate projection matrices.  Gather collisions intentionally produce
    gradient accumulation into the same source cells, the Dendro coupling mechanism.
    """

    def __init__(
        self,
        source_size: int | None = None,
        *,
        low_bit: bool = True,
        ternary_threshold: float = 0.7,
        init_std: float = 0.02,
        cache_derived_views: bool = False,
        cache_indices: bool = False,
        initial_source: torch.Tensor | None = None,
        logical_region_start: int = 0,
        logical_region_size: int | None = None,
        tensor_map: Mapping[str, Mapping[str, Any]] | None = None,
        aliases: Mapping[str, str] | None = None,
        exact_tied_logits: bool = False,
        freeze_named_tensors: bool = False,
    ) -> None:
        super().__init__()
        if initial_source is not None:
            if initial_source.ndim != 1:
                raise ValueError("initial_source must be a flat one-dimensional tensor")
            if not initial_source.is_floating_point():
                raise TypeError("initial_source must use a floating-point dtype")
            inferred_size = int(initial_source.numel())
            if source_size is not None and int(source_size) != inferred_size:
                raise ValueError(
                    f"source_size={source_size} does not match initial_source.numel()={inferred_size}"
                )
            source_size = inferred_size
        if source_size is None or int(source_size) <= 0:
            raise ValueError("source_size must be positive")
        self.source_size = int(source_size)
        self.logical_region_start = int(logical_region_start)
        self.logical_region_size = (
            self.source_size - self.logical_region_start
            if logical_region_size is None
            else int(logical_region_size)
        )
        if self.logical_region_start < 0:
            raise ValueError("logical_region_start cannot be negative")
        if self.logical_region_size <= 0:
            raise ValueError("logical_region_size must be positive")
        if self.logical_region_start + self.logical_region_size > self.source_size:
            raise ValueError("logical source region exceeds the physical source tensor")
        self.low_bit = bool(low_bit)
        self.ternary_threshold = float(ternary_threshold)
        self.cache_derived_views = bool(cache_derived_views)
        self.cache_indices = bool(cache_indices)
        if initial_source is None:
            source_tensor = torch.empty(self.source_size)
            nn.init.normal_(source_tensor, mean=0.0, std=float(init_std))
        else:
            source_tensor = initial_source.detach().contiguous()
        self.source = nn.Parameter(source_tensor)

        # A transplant tensor map gives symbolic names to non-overlapping slices of
        # the same physical parameter.  It does not register additional tensors.
        self._tensor_map: dict[str, dict[str, Any]] = {
            str(name): dict(record) for name, record in (tensor_map or {}).items()
        }
        self._aliases: dict[str, str] = {str(name): str(target) for name, target in (aliases or {}).items()}
        self.exact_tied_logits = bool(exact_tied_logits)
        self.freeze_named_tensors = bool(freeze_named_tensors)
        self._validate_tensor_map()

        # Python/runtime caches are deliberately not parameters or buffers.
        self._index_cache: dict[tuple[str, int, str], torch.Tensor] = {}
        self._eval_cache: dict[tuple[object, ...], torch.Tensor] = {}
        # One-forward derived views. Recurrent scans request the same logical
        # matrices once per token and depth; rebuilding them creates thousands of
        # identical autograd branches and dominated both runtime and VRAM. This
        # cache is non-persistent, contains no parameters/buffers, and is cleared
        # before the next grad-enabled model forward.
        self._forward_cache: dict[tuple[object, ...], torch.Tensor] = {}
        self._forward_cache_active = False
        self._forward_cache_source_version = self.source._version
        self._shape_book: dict[str, tuple[int, ...]] = {}
        self._source_version = self.source._version
        # Optional non-registered leaf used by the trainer for the logical Dendro
        # extension.  Keeping gradients on this small region avoids allocating and
        # repeatedly accumulating a dense gradient for a multi-gigabyte donor bank.
        self._training_extension: torch.Tensor | None = None
        # Full-weight training keeps donor gradients on the registered source and
        # uses only a small non-registered extension leaf. This permits separate
        # optimizer groups without cloning the multi-gigabyte donor prefix.
        self._training_donor: torch.Tensor | None = None
        self._joint_training_enabled = False
        self._source_requires_grad_before_extension = bool(self.source.requires_grad)

    _DTYPE_BY_NAME = {
        "float16": torch.float16,
        "half": torch.float16,
        "bfloat16": torch.bfloat16,
        "float32": torch.float32,
        "float": torch.float32,
        "float64": torch.float64,
    }

    def _validate_tensor_map(self) -> None:
        intervals: list[tuple[int, int, str]] = []
        for name, record in self._tensor_map.items():
            try:
                offset = int(record["offset"])
                shape = tuple(int(value) for value in record["shape"])
            except (KeyError, TypeError, ValueError) as exc:
                raise ValueError(f"invalid transplant tensor record for {name!r}: {record!r}") from exc
            if not shape or any(value <= 0 for value in shape):
                raise ValueError(f"transplant tensor {name!r} has invalid shape {shape}")
            numel = math.prod(shape)
            declared = int(record.get("numel", numel))
            if declared != numel:
                raise ValueError(f"transplant tensor {name!r} declares {declared} elements but shape has {numel}")
            if offset < 0 or offset + numel > self.source_size:
                raise ValueError(f"transplant tensor {name!r} exceeds the source tensor")
            record["offset"] = offset
            record["shape"] = list(shape)
            record["numel"] = numel
            intervals.append((offset, offset + numel, name))
        intervals.sort()
        for previous, current in zip(intervals, intervals[1:]):
            if current[0] < previous[1]:
                raise ValueError(
                    f"transplant tensors {previous[2]!r} and {current[2]!r} overlap in the single source"
                )
        for alias, target in self._aliases.items():
            if target not in self._tensor_map:
                raise ValueError(f"source alias {alias!r} points to unknown transplant tensor {target!r}")

    @property
    def tensor_map(self) -> dict[str, dict[str, Any]]:
        return {name: dict(record) for name, record in self._tensor_map.items()}

    @property
    def aliases(self) -> dict[str, str]:
        return dict(self._aliases)

    def has_named_tensor(self, name: str) -> bool:
        return str(name) in self._tensor_map

    def named_tensor(
        self,
        name: str,
        *,
        dtype: torch.dtype | None = None,
    ) -> torch.Tensor:
        """Return an exact symbolic view from the one physical source parameter.

        The checkpoint manifest records the donor dtype.  FP32 source storage can
        therefore reproduce both BF16 and FP32 donor tensors exactly; BF16 storage
        intentionally rounds the small FP32 state subset and is marked non-exact by
        the transplant manifest.
        """

        key = str(name)
        if key not in self._tensor_map:
            raise KeyError(f"unknown transplant tensor {key!r}")
        record = self._tensor_map[key]
        target_dtype = dtype
        if target_dtype is None:
            dtype_name = str(record.get("dtype", "")).replace("torch.", "").lower()
            target_dtype = self._DTYPE_BY_NAME.get(dtype_name)
        cache_key = (
            "named_tensor",
            key,
            target_dtype,
            self.freeze_named_tensors,
            self.source._version,
            None if self._training_donor is None else self._training_donor._version,
            None if self._training_extension is None else self._training_extension._version,
            self.source.device,
            self.source.dtype,
        )
        cached = self._forward_cached(cache_key)
        if cached is not None:
            return cached
        offset = int(record["offset"])
        numel = int(record["numel"])
        shape = tuple(int(value) for value in record["shape"])
        end = offset + numel
        logical_end = self.logical_region_start + self.logical_region_size
        if self._training_donor is not None and end <= self.logical_region_start:
            tensor = self._training_donor.narrow(0, offset, numel).reshape(shape)
        elif (
            self._training_extension is not None
            and offset >= self.logical_region_start
            and end <= logical_end
        ):
            tensor = self._training_extension.narrow(
                0, offset - self.logical_region_start, numel
            ).reshape(shape)
        elif self._training_donor is not None and offset < self.logical_region_start < end:
            raise RuntimeError(
                f"transplant tensor {key!r} crosses the donor/extension training boundary"
            )
        else:
            tensor = self.source.narrow(0, offset, numel).reshape(shape)
        # Freeze only the named donor coordinates.  Dendro-derived primitives still
        # read the trainable extension region of this same physical Parameter.
        if self.freeze_named_tensors:
            tensor = tensor.detach()
        if target_dtype is not None and tensor.dtype != target_dtype:
            tensor = tensor.to(dtype=target_dtype)
        return self._store_forward_cached(cache_key, tensor)

    def resolve_alias(self, logical_name: str) -> str | None:
        return self._aliases.get(str(logical_name))

    @staticmethod
    @lru_cache(maxsize=32_768)
    def _digest(name: str) -> tuple[int, int, int]:
        digest = hashlib.blake2b(name.encode("utf-8"), digest_size=24).digest()
        return (
            int.from_bytes(digest[0:8], "little"),
            int.from_bytes(digest[8:16], "little"),
            int.from_bytes(digest[16:24], "little"),
        )

    def _clear_if_changed(self) -> None:
        if self._source_version != self.source._version:
            self._eval_cache.clear()
            self._source_version = self.source._version

    def clear_runtime_cache(self) -> None:
        self._index_cache.clear()
        self._eval_cache.clear()
        self._forward_cache.clear()

    def begin_forward_cache(self) -> None:
        """Start an autograd-safe cache shared by every recurrent depth.

        Training graphs may never be reused across microbatches, so a
        grad-enabled call always begins empty. Inference can retain the same
        derived tensors across cached decode calls until the source changes.
        """

        changed = self._forward_cache_source_version != self.source._version
        if torch.is_grad_enabled() or changed:
            self._forward_cache.clear()
        self._forward_cache_source_version = self.source._version
        self._forward_cache_active = True

    def _forward_cached(self, key: tuple[object, ...]) -> torch.Tensor | None:
        if not self._forward_cache_active:
            return None
        return self._forward_cache.get(key)

    def _store_forward_cached(self, key: tuple[object, ...], value: torch.Tensor) -> torch.Tensor:
        if self._forward_cache_active:
            self._forward_cache[key] = value
        return value

    @property
    def training_extension_active(self) -> bool:
        return self._training_extension is not None

    @property
    def training_donor_active(self) -> bool:
        return bool(self._joint_training_enabled)

    @property
    def joint_training_active(self) -> bool:
        return bool(self._joint_training_enabled and self._training_extension is not None)

    def enable_training_extension(
        self,
        *,
        master_dtype: torch.dtype | None = None,
    ) -> torch.Tensor:
        """Return a trainable leaf for only the reserved logical source region.

        The leaf is deliberately neither a parameter nor a buffer.  Checkpoints
        therefore retain exactly one registered physical source parameter.  The
        trainer commits the leaf back to that parameter after optimizer steps and
        before serialization.
        """

        if self._training_extension is not None:
            return self._training_extension
        self._source_requires_grad_before_extension = bool(self.source.requires_grad)
        self.source.grad = None
        self.source.requires_grad_(False)
        extension = (
            self.source.detach()
            .narrow(0, self.logical_region_start, self.logical_region_size)
            .to(dtype=self.source.dtype if master_dtype is None else master_dtype)
            .clone()
            .requires_grad_(True)
        )
        self._training_extension = extension
        self.clear_runtime_cache()
        return extension

    def enable_joint_training(
        self,
        *,
        extension_master_dtype: torch.dtype | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Return the registered donor source and a non-registered extension leaf.

        The physical ``source`` remains the model's only registered parameter.
        Named transplanted tensors read from ``source`` while Dendro primitives
        read from ``extension``; :meth:`commit_joint_training` copies the small
        extension leaf back into the source at optimizer/checkpoint boundaries.
        """

        if self.joint_training_active:
            assert self._training_extension is not None
            return self.source, self._training_extension
        if self._joint_training_enabled or self._training_extension is not None:
            self.disable_joint_training(commit=True)
        logical_end = self.logical_region_start + self.logical_region_size
        if logical_end != self.source_size:
            raise ValueError(
                "joint training requires the Dendro extension to cover the physical source suffix"
            )
        if self.logical_region_start <= 0:
            raise ValueError("joint training requires a non-empty transplanted donor prefix")
        self._source_requires_grad_before_extension = bool(self.source.requires_grad)
        self.source.grad = None
        self.source.requires_grad_(True)
        extension = (
            self.source.detach()
            .narrow(0, self.logical_region_start, self.logical_region_size)
            .to(
                dtype=(
                    self.source.dtype
                    if extension_master_dtype is None
                    else extension_master_dtype
                )
            )
            .clone()
            .requires_grad_(True)
        )
        self._training_donor = None
        self._joint_training_enabled = True
        self._training_extension = extension
        self.clear_runtime_cache()
        return self.source, extension

    @torch.no_grad()
    def commit_training_extension(self) -> None:
        if self._training_extension is None:
            return
        target = self.source.narrow(0, self.logical_region_start, self.logical_region_size)
        target.copy_(self._training_extension.detach().to(device=target.device, dtype=target.dtype))
        self.clear_runtime_cache()

    @torch.no_grad()
    def commit_joint_training(self) -> None:
        if self._training_extension is not None:
            extension_target = self.source.narrow(
                0, self.logical_region_start, self.logical_region_size
            )
            extension_target.copy_(
                self._training_extension.detach().to(
                    device=extension_target.device, dtype=extension_target.dtype
                )
            )
        self.clear_runtime_cache()

    def disable_training_extension(self, *, commit: bool = True) -> None:
        if self._training_extension is None:
            return
        if self._joint_training_enabled:
            self.disable_joint_training(commit=commit)
            return
        if commit:
            self.commit_training_extension()
        self._training_extension = None
        self.source.requires_grad_(self._source_requires_grad_before_extension)
        self.clear_runtime_cache()

    def disable_joint_training(self, *, commit: bool = True) -> None:
        if not self._joint_training_enabled and self._training_extension is None:
            return
        if commit:
            self.commit_joint_training()
        self._training_donor = None
        self._joint_training_enabled = False
        self._training_extension = None
        self.source.requires_grad_(self._source_requires_grad_before_extension)
        self.clear_runtime_cache()

    def _can_cache_values(self) -> bool:
        return (
            self._training_donor is None
            and self._training_extension is None
            and self.cache_derived_views
            and not self.training
            and not torch.is_grad_enabled()
        )

    def _contiguous_offset(self, name: str, numel: int) -> int | None:
        if numel > self.logical_region_size:
            return None
        h0, _h1, _h2 = self._digest(name)
        local = h0 % max(1, self.logical_region_size - numel + 1)
        return self.logical_region_start + local

    def _indices(self, name: str, numel: int, device: torch.device) -> torch.Tensor:
        key = (name, int(numel), str(device))
        if self.cache_indices:
            cached = self._index_cache.get(key)
            if cached is not None:
                return cached
        contiguous = self._contiguous_offset(name, numel)
        if contiguous is not None:
            indices = torch.arange(contiguous, contiguous + numel, device=device, dtype=torch.long)
        else:
            h0, h1, h2 = self._digest(name)
            offset = h0 % self.logical_region_size
            # Oversized logical tensors wrap and couple repeatedly through the reserved
            # logical region.  This lets a lossless donor bank occupy the beginning of
            # the physical source without Dendro-derived views overwriting it.
            stride = (h1 % max(1, self.logical_region_size - 1)) + 1
            if stride % 2 == 0 and self.logical_region_size > 1:
                stride += 1
            phase = (h2 % 97) + 1
            positions = torch.arange(numel, device=device, dtype=torch.long)
            local = torch.remainder(
                offset + positions * stride + (positions // phase) ** 2,
                self.logical_region_size,
            )
            indices = local + self.logical_region_start
        if self.cache_indices:
            self._index_cache[key] = indices
        return indices

    def primitive(self, name: str, shape: Sequence[int], *, scale: float = 1.0) -> torch.Tensor:
        checked = tuple(int(dim) for dim in shape)
        if not checked or any(dim <= 0 for dim in checked):
            raise ValueError(f"primitive {name!r} requires a positive shape, got {checked}")
        previous = self._shape_book.get(name)
        if previous is not None and previous != checked:
            raise ValueError(f"primitive {name!r} requested as both {previous} and {checked}")
        self._shape_book[name] = checked
        numel = math.prod(checked)
        alias = self.resolve_alias(name)
        if alias is not None:
            result = self.named_tensor(alias)
            if tuple(result.shape) != checked:
                raise ValueError(
                    f"source alias {name!r} -> {alias!r} has shape {tuple(result.shape)}, expected {checked}"
                )
            return result if scale == 1.0 else result * float(scale)
        cache_key = ("primitive", name, checked, float(scale), self.source._version, self.source.device, self.source.dtype)
        forward_cached = self._forward_cached(cache_key)
        if forward_cached is not None:
            return forward_cached
        if self._can_cache_values():
            self._clear_if_changed()
            cached = self._eval_cache.get(cache_key)
            if cached is not None:
                return cached
        contiguous = self._contiguous_offset(name, numel)
        extension = self._training_extension
        if extension is not None and contiguous is not None:
            result = extension.narrow(
                0, contiguous - self.logical_region_start, numel
            ).reshape(checked)
        elif extension is not None:
            indices = self._indices(name, numel, extension.device) - self.logical_region_start
            result = extension.index_select(0, indices).reshape(checked)
        elif contiguous is not None:
            result = self.source.narrow(0, contiguous, numel).reshape(checked)
        else:
            indices = self._indices(name, numel, self.source.device)
            result = self.source.index_select(0, indices).reshape(checked)
        # Optimizers need an FP32 master at the tiny learning rates used for
        # routing calibration; applying every step directly to a BF16 leaf
        # rounds most updates to zero.  Cast the derived view back to the
        # physical source dtype for byte-compatible forward numerics.  The cast
        # remains differentiable, so gradients accumulate on the FP32 master.
        if extension is not None and result.dtype != self.source.dtype:
            result = result.to(dtype=self.source.dtype)
        if scale != 1.0:
            result = result * float(scale)
        if self._can_cache_values():
            self._eval_cache[cache_key] = result
        return self._store_forward_cached(cache_key, result)

    def ternary(self, tensor: torch.Tensor) -> torch.Tensor:
        """Straight-through 1.58-bit {-scale, 0, +scale} quantization."""

        if tensor.ndim >= 2:
            reduce_dims = tuple(range(1, tensor.ndim))
        else:
            reduce_dims = (0,)
        scale = tensor.detach().abs().mean(dim=reduce_dims, keepdim=True).clamp_min(1e-8)
        threshold = self.ternary_threshold * scale
        quantized = torch.where(
            tensor > threshold,
            scale,
            torch.where(tensor < -threshold, -scale, torch.zeros_like(tensor)),
        )
        return tensor + (quantized - tensor).detach()

    def weight(
        self,
        name: str,
        out_features: int,
        in_features: int,
        *,
        low_bit: bool | None = None,
    ) -> torch.Tensor:
        use_low_bit = self.low_bit if low_bit is None else bool(low_bit)
        key = (
            "weight",
            name,
            int(out_features),
            int(in_features),
            use_low_bit,
            self.source._version,
            self.source.device,
            self.source.dtype,
        )
        forward_cached = self._forward_cached(key)
        if forward_cached is not None:
            return forward_cached
        if self._can_cache_values():
            self._clear_if_changed()
            cached = self._eval_cache.get(key)
            if cached is not None:
                return cached
        raw = self.primitive(f"{name}/weight", (int(out_features), int(in_features)))
        weight = raw / math.sqrt(max(1, int(in_features)))
        if use_low_bit:
            weight = self.ternary(weight)
        if self._can_cache_values():
            self._eval_cache[key] = weight
        return self._store_forward_cached(key, weight)

    def bias(self, name: str, features: int, *, magnitude: float = 0.01) -> torch.Tensor:
        key = ("bias", name, int(features), float(magnitude), self.source._version, self.source.device, self.source.dtype)
        cached = self._forward_cached(key)
        if cached is not None:
            return cached
        return self._store_forward_cached(
            key,
            self.primitive(f"{name}/bias", (int(features),), scale=float(magnitude)),
        )

    def project(
        self,
        x: torch.Tensor,
        name: str,
        out_features: int,
        *,
        bias: bool = True,
        low_bit: bool | None = None,
    ) -> torch.Tensor:
        weight = self.weight(name, int(out_features), int(x.shape[-1]), low_bit=low_bit)
        bias_tensor = self.bias(name, int(out_features)) if bias else None
        return F.linear(x, weight, bias_tensor)

    def project_many(
        self,
        x: torch.Tensor,
        specs: Iterable[tuple[str, int, bool | None]],
    ) -> tuple[torch.Tensor, ...]:
        checked = tuple((str(name), int(size), low_bit) for name, size, low_bit in specs)
        if not checked or any(size <= 0 for _name, size, _low_bit in checked):
            raise ValueError("project_many requires one or more positive output sizes")
        weights = [self.weight(name, size, int(x.shape[-1]), low_bit=low_bit) for name, size, low_bit in checked]
        biases = [self.bias(name, size) for name, size, _low_bit in checked]
        merged = F.linear(x, torch.cat(weights, dim=0), torch.cat(biases, dim=0))
        return merged.split([size for _name, size, _low_bit in checked], dim=-1)

    def embedding(
        self,
        ids: torch.Tensor,
        name: str,
        num_embeddings: int,
        embedding_dim: int,
        *,
        low_bit: bool = False,
    ) -> torch.Tensor:
        table = self.primitive(f"{name}/embedding", (int(num_embeddings), int(embedding_dim)))
        if low_bit:
            table = self.ternary(table)
        return F.embedding(ids, table)

    def rms_norm(
        self,
        x: torch.Tensor,
        name: str,
        *,
        eps: float = 1e-5,
        shared_scale: torch.Tensor | None = None,
    ) -> torch.Tensor:
        normalized = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps).to(x.dtype)
        scale = shared_scale
        if scale is None:
            scale = 1.0 + 0.05 * torch.tanh(self.primitive(f"{name}/scale", (int(x.shape[-1]),)))
        return normalized * scale

    def aligned_qkv_norm(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        name: str,
        *,
        eps: float = 1e-5,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Normalize Q/K/V with the exact same source-derived head scale."""

        head_dim = int(q.shape[-1])
        shared_scale = 1.0 + 0.05 * torch.tanh(self.primitive(f"{name}/shared_scale", (head_dim,)))
        return (
            self.rms_norm(q, f"{name}/q", eps=eps, shared_scale=shared_scale),
            self.rms_norm(k, f"{name}/k", eps=eps, shared_scale=shared_scale),
            self.rms_norm(v, f"{name}/v", eps=eps, shared_scale=shared_scale),
        )

    def gate(self, x: torch.Tensor, name: str, out_features: int = 1) -> torch.Tensor:
        return torch.sigmoid(self.project(x, name, out_features, low_bit=False))

    def tied_logits(self, hidden: torch.Tensor, *, vocab_size: int, name: str = "token") -> torch.Tensor:
        logical_name = f"{name}/embedding"
        table = self.primitive(logical_name, (int(vocab_size), int(hidden.shape[-1])))
        if self.exact_tied_logits and self.resolve_alias(logical_name) is not None:
            return F.linear(hidden, table.to(dtype=hidden.dtype))
        return F.linear(hidden, table / math.sqrt(max(1, int(hidden.shape[-1]))), self.bias("lm_head", vocab_size))

    def tied_selected_logits(
        self,
        hidden: torch.Tensor,
        *,
        vocab_indices: torch.Tensor,
        vocab_size: int,
        name: str = "token",
    ) -> torch.Tensor:
        """Project a training batch onto selected tied vocabulary rows.

        This is an exact row selection from the same embedding/source tensor, not
        a second head or a new parameter.  It lets sampled-softmax training avoid
        materializing a 248k-wide logit tensor while ordinary inference retains
        the complete vocabulary projection.
        """

        indices = vocab_indices.to(device=hidden.device, dtype=torch.long).flatten()
        if indices.numel() == 0:
            raise ValueError("vocab_indices cannot be empty")
        if int(indices.min().item()) < 0 or int(indices.max().item()) >= int(vocab_size):
            raise ValueError("vocab_indices contains an out-of-vocabulary token")
        logical_name = f"{name}/embedding"
        table = self.primitive(logical_name, (int(vocab_size), int(hidden.shape[-1])))
        selected = table.index_select(0, indices)
        if self.exact_tied_logits and self.resolve_alias(logical_name) is not None:
            return F.linear(hidden, selected.to(dtype=hidden.dtype))
        bias = self.bias("lm_head", vocab_size).index_select(0, indices)
        return F.linear(
            hidden,
            selected / math.sqrt(max(1, int(hidden.shape[-1]))),
            bias,
        )

    @torch.no_grad()
    def write_primitive(self, name: str, value: torch.Tensor) -> None:
        """Project an external logical tensor back into the one source by averaging collisions."""

        shape = tuple(int(dim) for dim in value.shape)
        logical_name = name
        if self.resolve_alias(logical_name) is not None:
            raise RuntimeError(
                f"logical primitive {logical_name!r} aliases the immutable transplant donor bank; "
                "write an extension primitive or explicitly rebuild the transplant checkpoint"
            )
        previous = self._shape_book.get(logical_name)
        if previous is not None and previous != shape:
            self._shape_book.pop(logical_name, None)
        self._shape_book[logical_name] = shape
        indices = self._indices(logical_name, value.numel(), self.source.device)
        flat = value.detach().to(device=self.source.device, dtype=self.source.dtype).reshape(-1)
        # Allocate only the active logical region.  Exact-transplant source banks can
        # contain hundreds of millions of donor elements, while Dendro's extension
        # region is intentionally much smaller.
        local_indices = indices - self.logical_region_start
        sums = torch.zeros(
            self.logical_region_size,
            device=self.source.device,
            dtype=self.source.dtype,
        )
        counts = torch.zeros_like(sums)
        sums.scatter_add_(0, local_indices, flat)
        counts.scatter_add_(0, local_indices, torch.ones_like(flat))
        touched = counts > 0
        target = self.source.data.narrow(0, self.logical_region_start, self.logical_region_size)
        target[touched] = sums[touched] / counts[touched]
        self.clear_runtime_cache()

    def forget_primitive_shape(self, name: str) -> None:
        """Allow a logical view such as a resized vocabulary to request a new shape."""

        self._shape_book.pop(name, None)
        self._shape_book.pop(f"{name}/embedding", None)
        self.clear_runtime_cache()

    def pack_primitive(self, name: str, shape: Sequence[int]) -> PackedTernaryTensor:
        with torch.no_grad():
            tensor = self.primitive(name, shape)
            if tensor.ndim >= 2:
                reduce_dims = tuple(range(1, tensor.ndim))
            else:
                reduce_dims = (0,)
            scale = tensor.abs().mean(dim=reduce_dims, keepdim=True).clamp_min(1e-8)
            threshold = self.ternary_threshold * scale
            codes = torch.where(tensor > threshold, 1, torch.where(tensor < -threshold, 2, 0)).to(torch.uint8).flatten()
            pad = (-codes.numel()) % 4
            if pad:
                codes = F.pad(codes, (0, pad))
            codes = codes.reshape(-1, 4)
            shifts = torch.tensor([0, 2, 4, 6], device=codes.device, dtype=torch.uint8)
            packed = torch.sum(codes << shifts, dim=-1).to(torch.uint8)
            return PackedTernaryTensor(packed=packed, scale=scale.detach(), shape=tuple(shape), numel=math.prod(shape))

    def audit(self) -> dict[str, int | bool]:
        parameters = list(self.parameters())
        return {
            "registered_parameter_tensors": len(parameters),
            "registered_parameter_elements": sum(parameter.numel() for parameter in parameters),
            "source_elements": self.source.numel(),
            "logical_region_start": self.logical_region_start,
            "logical_region_elements": self.logical_region_size,
            "literal_single_source": len(parameters) == 1 and parameters[0] is self.source,
            "low_bit_views": self.low_bit,
            "logical_primitives_requested": len(self._shape_book),
            "cached_indices": len(self._index_cache),
            "cached_eval_values": len(self._eval_cache),
            "cached_forward_values": len(self._forward_cache),
            "persistent_derived_tensor_bytes": sum(
                tensor.numel() * tensor.element_size() for tensor in self._index_cache.values()
            ) + sum(tensor.numel() * tensor.element_size() for tensor in self._eval_cache.values()),
            "cache_derived_views": self.cache_derived_views,
            "cache_indices": self.cache_indices,
            "transplant_named_tensors": len(self._tensor_map),
            "transplant_aliases": len(self._aliases),
            "exact_tied_logits": self.exact_tied_logits,
            "freeze_named_tensors": self.freeze_named_tensors,
            "training_donor_active": self.training_donor_active,
            "training_donor_elements": (
                self.logical_region_start if self._joint_training_enabled else 0
            ),
            "training_extension_active": self.training_extension_active,
            "training_extension_elements": (
                0 if self._training_extension is None else self._training_extension.numel()
            ),
        }