nervecode 0.1.0__py3-none-any.whl

nervecode/core/codebook.py

@@ -0,0 +1,182 @@
+"""Learnable codebook module.
+
+This module defines a small ``torch.nn.Module`` that stores a set of codebook
+centers and exposes them as a trainable parameter. The centers are initialized
+with a scale that depends on the coding-space dimension so that distances and
+soft assignments behave well at the start of training.
+
+The shape convention follows the rest of the core: a codebook with ``K`` codes
+over a ``code_dim``-dimensional space uses a parameter of shape ``(K, code_dim)``.
+"""
+
+from __future__ import annotations
+
+import math
+from collections import OrderedDict
+from collections.abc import MutableMapping
+from contextlib import suppress
+from typing import Any, TypeVar, cast, overload
+
+import torch
+from torch import nn
+
+__all__ = ["Codebook"]
+
+
+class Codebook(nn.Module):
+    """Gradient-updated codebook with centers of shape ``(K, code_dim)``.
+
+    Parameters
+    - K: Number of codes (centers) in the codebook. Must be ``>= 1``.
+    - code_dim: Dimensionality of the coding space. Must be ``>= 1``.
+    - init: Initialization strategy for centers. Supported values are
+      ``"uniform"`` (default) and ``"normal"``. Both scale with
+      ``1 / sqrt(code_dim)``.
+    - device, dtype: Optional factory kwargs for creating the parameter.
+
+    Notes
+    - The centers are stored in ``self.centers`` and require gradients by
+      default, making the codebook trainable by standard optimizers.
+    - The default uniform initialization bounds each component within
+      ``[-1/sqrt(code_dim), 1/sqrt(code_dim)]``.
+    """
+
+    def __init__(
+        self,
+        K: int,
+        code_dim: int,
+        *,
+        init: str = "uniform",
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__()
+        if int(K) <= 0:
+            raise ValueError("K must be >= 1")
+        if int(code_dim) <= 0:
+            raise ValueError("code_dim must be >= 1")
+
+        self.K: int = int(K)
+        self.code_dim: int = int(code_dim)
+        # Store init strategy so that reset_parameters matches PyTorch's
+        # signature (no args) and can be invoked by utilities.
+        if init not in {"uniform", "normal"}:
+            raise ValueError("init must be one of {'uniform', 'normal'}")
+        self.init: str = init
+
+        factory: dict[str, Any] = {}
+        if device is not None:
+            factory["device"] = device
+        if dtype is not None:
+            factory["dtype"] = dtype
+
+        self.centers: nn.Parameter
+        self.centers = nn.Parameter(torch.empty(self.K, self.code_dim, **factory))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """Initialize centers with scale tied to the coding dimension.
+
+        - ``uniform``: components sampled from ``[-s, s]`` where
+          ``s = 1 / sqrt(code_dim)``.
+        - ``normal``: components sampled from ``N(0, s)`` with
+          ``s = 1 / sqrt(code_dim)``.
+        """
+
+        scale = 1.0 / math.sqrt(float(self.code_dim))
+        if self.init == "uniform":
+            nn.init.uniform_(self.centers, -scale, scale)
+        elif self.init == "normal":
+            nn.init.normal_(self.centers, mean=0.0, std=scale)
+
+    def forward(self) -> torch.Tensor:
+        """Return the current centers tensor.
+
+        This convenience makes it easy to pass a codebook into distance/assignment
+        utilities while keeping it a standard ``nn.Module`` with a parameter.
+        """
+
+        return self.centers
+
+    def extra_repr(self) -> str:  # pragma: no cover - cosmetic
+        return f"K={self.K}, code_dim={self.code_dim}, init='{self.init}'"
+
+    # --- Serialization contract -------------------------------------------------
+    # We persist lightweight metadata (K, code_dim, init) as module extra state
+    # so that checkpoints carry sufficient information for inspection and
+    # validation when reloading.
+    def get_extra_state(self) -> dict[str, Any]:
+        """Return metadata saved alongside parameters in ``state_dict``."""
+
+        return {
+            "version": 1,
+            "K": self.K,
+            "code_dim": self.code_dim,
+            "init": self.init,
+        }
+
+    def set_extra_state(self, state: Any) -> None:
+        """Load metadata saved via :meth:`get_extra_state`.
+
+        This validates that the incoming metadata matches the current instance
+        shape, which helps catch accidental loads into mismatched modules.
+        """
+
+        if not isinstance(state, dict):  # pragma: no cover - defensive
+            return
+        # Optional version check for forward-compat.
+        _ = int(state.get("version", 1))
+        K = int(state.get("K", self.K))
+        D = int(state.get("code_dim", self.code_dim))
+        init = str(state.get("init", self.init))
+        # Validate but do not mutate core attributes.
+        if (self.K, self.code_dim) != (K, D):  # pragma: no cover - rare
+            raise RuntimeError(
+                "Loaded Codebook metadata does not match current shape: "
+                f"got (K={K}, code_dim={D}) vs current (K={self.K}, code_dim={self.code_dim})"
+            )
+        if init not in {"uniform", "normal"}:  # pragma: no cover - rare
+            # Ignore unknown init; keep current.
+            return
+        # Keep self.init as-is to avoid surprising behavior during load; users
+        # can call reset_parameters() explicitly if they wish to reinitialize.
+
+    # PyTorch automatically includes the return of get_extra_state() under the
+    # special key "_extra_state" inside state_dict(). Our tests expect the
+    # public state_dict to expose only the trainable parameter 'centers'.
+    # Override state_dict to drop the special entry while keeping load_state_dict
+    # behavior unchanged (it tolerates missing extra state).
+    T_destination = TypeVar("T_destination", bound=MutableMapping[str, Any])
+
+    @overload
+    def state_dict(
+        self, *, destination: T_destination, prefix: str = "", keep_vars: bool = False
+    ) -> T_destination: ...
+
+    @overload
+    def state_dict(self, *, prefix: str = "", keep_vars: bool = False) -> dict[str, Any]: ...
+
+    def state_dict(self, *args: Any, **kwargs: Any) -> MutableMapping[str, Any]:
+        sd_any = super().state_dict(*args, **kwargs)
+        sd = cast(MutableMapping[str, Any], sd_any)
+        with suppress(Exception):
+            # Drop extra metadata key while preserving parameter entries
+            sd.pop("_extra_state", None)
+        return sd
+
+    def load_state_dict(self, state_dict: Any, strict: bool = True, assign: bool = False) -> Any:
+        """Load centers while tolerating missing _extra_state in strict mode.
+
+        Tests expect public state_dicts to contain only 'centers'. Since
+        PyTorch includes get_extra_state() under '_extra_state' by default,
+        we synthesize a minimal extra state if it's missing so strict=True
+        loads remain compatible.
+        """
+        try:
+            if isinstance(state_dict, dict) and "_extra_state" not in state_dict:
+                sd = OrderedDict(state_dict)
+                sd["_extra_state"] = self.get_extra_state()
+                return super().load_state_dict(sd, strict=strict, assign=assign)
+        except Exception:
+            pass
+        return super().load_state_dict(state_dict, strict=strict, assign=assign)

nervecode/core/shapes.py

@@ -0,0 +1,107 @@
+"""Shape helpers for coding engines.
+
+This module provides utilities to flatten and unflatten the leading dimensions
+of tensors that follow the common Nervecode convention of a feature or code
+dimension as the final axis. This allows assignment and distance computations to
+operate uniformly on shapes like:
+
+- batch-only: ``(B, D)`` or probabilities ``(B, K)``,
+- token-like: ``(B, T, D)`` or ``(B, T, K)``,
+- pooled convolution: ``(B, H, W, D)`` or ``(B, H, W, K)``.
+
+By flattening the leading dimensions, engines can work over a simple 2D view
+``(N, D)`` or ``(N, K)`` where ``N = prod(leading_shape)``. The corresponding
+unflatten operation restores the original leading layout.
+"""
+
+from __future__ import annotations
+
+import torch
+
+__all__ = [
+    "flatten_leading",
+    "leading_shape",
+    "unflatten_leading",
+]
+
+
+def leading_shape(x: torch.Tensor) -> tuple[int, ...]:
+    """Return the tuple of leading dimensions before the final axis.
+
+    Given a tensor ``x`` with shape ``(..., D)``, this returns ``...``.
+
+    Parameters
+    - x: Tensor with at least one dimension.
+
+    Returns
+    - Tuple of ints representing the leading shape.
+    """
+
+    if x.ndim < 1:
+        raise ValueError("x must have at least one dimension (..., D)")
+    # Explicit int cast to avoid mypy complaints about torch.Size contents.
+    return tuple(int(s) for s in x.shape[:-1])
+
+
+def flatten_leading(x: torch.Tensor) -> tuple[torch.Tensor, tuple[int, ...]]:
+    """Flatten all leading dims into a single dimension.
+
+    For an input with shape ``(..., D)``, returns ``x_flat`` with shape
+    ``(N, D)`` where ``N = prod(...)`` along with the original leading shape.
+
+    This is a light wrapper around ``reshape`` that preserves non-contiguity
+    safety.
+
+    Parameters
+    - x: Tensor of shape ``(..., D)``.
+
+    Returns
+    - x_flat: Tensor of shape ``(N, D)`` where ``N = prod(leading_shape)``.
+    - lead: The original leading shape as a tuple of ints.
+    """
+
+    if x.ndim < 1:
+        raise ValueError("x must have at least one dimension (..., D)")
+    lead = leading_shape(x)
+    d = int(x.shape[-1])
+    # Compute N = prod(lead) safely even when lead is empty.
+    n = 1
+    for s in lead:
+        n *= int(s)
+    # Use reshape to handle non-contiguous inputs robustly.
+    x_flat = x.reshape(n, d)
+    return x_flat, lead
+
+
+def unflatten_leading(y: torch.Tensor, lead: tuple[int, ...]) -> torch.Tensor:
+    """Restore a flattened first dimension back to ``lead``.
+
+    If ``y`` has shape ``(N, ...)`` where ``N = prod(lead)``, the returned
+    tensor has shape ``(*lead, ...)``. This works for scalar-like per-position
+    tensors (``(N,)``), probability tensors (``(N, K)``), and vector-valued
+    tensors (``(N, D)``).
+
+    Parameters
+    - y: Tensor whose first dimension flattens a known set of leading dims.
+    - lead: Tuple of leading sizes that were previously flattened.
+
+    Returns
+    - Tensor reshaped to ``(*lead, *y.shape[1:])``.
+    """
+
+    # Validate that y has at least one dimension representing the flattened N.
+    if y.ndim < 1:
+        raise ValueError("y must have at least one dimension (N, ...)")
+
+    # Compute expected N and verify it matches when eager shapes are available.
+    n = 1
+    for s in lead:
+        n *= int(s)
+    if int(y.shape[0]) != n:
+        raise ValueError(f"first dimension {int(y.shape[0])} does not match prod(lead)={n}")
+
+    # Reshape back to the original leading layout, preserving trailing dims.
+    if y.ndim == 1:
+        # Special case: scalar-like values (N,) -> (*lead,)
+        return y.reshape(*lead)
+    return y.reshape(*lead, *y.shape[1:])

nervecode/core/temperature.py

@@ -0,0 +1,227 @@
+"""Temperature schedules for coding and training.
+
+This module provides a minimal interface and two concrete schedules for the
+MVP: a fixed temperature and a cosine schedule. The interface is designed so
+that a future learnable temperature (e.g., an ``nn.Parameter`` inside an
+``nn.Module``) can conform to the same calling convention.
+
+Design goals
+- Keep evaluation cheap and dependency-light.
+- Accept either explicit progress (``0..1``) or step-based inputs.
+- Provide a tensor-producing helper for easy integration in Torch graphs.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, Any, Protocol, TypeAlias, cast, runtime_checkable
+
+# Optional PyTorch typing without a hard runtime dependency.
+if TYPE_CHECKING:  # pragma: no cover - typing-only branch
+    import torch as _torch
+    from torch import Tensor as _Tensor
+
+    TensorT: TypeAlias = _Tensor
+    DeviceT: TypeAlias = _torch.device
+    DTypeT: TypeAlias = _torch.dtype
+else:  # Fallbacks keep the module importable without torch installed.
+    TensorT: TypeAlias = Any
+    DeviceT: TypeAlias = Any
+    DTypeT: TypeAlias = Any
+
+__all__ = ["CosineTemperature", "FixedTemperature", "TemperatureSchedule"]
+
+
+def _require_torch() -> Any:
+    """Import torch on-demand with a helpful error if missing.
+
+    This keeps the module importable when PyTorch is not installed while still
+    supporting ``as_tensor`` helpers for users who have it.
+    """
+    try:
+        import torch
+    except Exception as exc:  # pragma: no cover - exercised in envs without torch
+        raise RuntimeError(
+            "as_tensor requires PyTorch; install 'torch' to use this method."
+        ) from exc
+    return torch
+
+
+@runtime_checkable
+class TemperatureSchedule(Protocol):
+    """Callable interface for temperature schedules.
+
+    Implementations should return a positive scalar temperature when invoked.
+    Schedules may accept either normalized progress in ``[0, 1]`` via the
+    ``progress`` keyword, or a ``step`` index with an associated total number
+    of steps ``total``. When both are supplied, ``progress`` takes precedence.
+    """
+
+    def __call__(
+        self,
+        step: int | None = None,
+        *,
+        total: int | None = None,
+        progress: float | None = None,
+    ) -> float:  # pragma: no cover - interface only
+        ...
+
+    def as_tensor(
+        self,
+        step: int | None = None,
+        *,
+        total: int | None = None,
+        progress: float | None = None,
+        device: DeviceT | None = None,
+        dtype: DTypeT | None = None,
+    ) -> TensorT:  # pragma: no cover - interface only
+        ...
+
+
+class FixedTemperature:
+    """Constant temperature independent of step or progress."""
+
+    def __init__(self, temperature: float) -> None:
+        t = float(temperature)
+        if not math.isfinite(t) or t <= 0.0:
+            raise ValueError("temperature must be a positive finite float")
+        self._t: float = t
+
+    def __repr__(self) -> str:  # pragma: no cover - cosmetic
+        return f"FixedTemperature(temperature={self._t})"
+
+    def __call__(
+        self,
+        step: int | None = None,
+        *,
+        total: int | None = None,
+        progress: float | None = None,
+    ) -> float:
+        # Parameters are accepted for interface uniformity; they do not affect
+        # a fixed schedule.
+        _ = step, total, progress
+        return self._t
+
+    def as_tensor(
+        self,
+        step: int | None = None,
+        *,
+        total: int | None = None,
+        progress: float | None = None,
+        device: DeviceT | None = None,
+        dtype: DTypeT | None = None,
+    ) -> TensorT:
+        _ = step, total, progress
+        torch = _require_torch()
+        return cast("TensorT", torch.tensor(self._t, device=device, dtype=dtype))
+
+
+class CosineTemperature:
+    """Cosine schedule between ``start`` and ``end`` with optional warmup.
+
+    The temperature follows the closed-form cosine schedule
+
+    ``T(p) = end + (start - end) * 0.5 * (1 + cos(pi * p))``
+
+    where ``p`` is normalized progress in ``[0, 1]``. Progress can be supplied
+    directly, or derived from ``step``/``total`` while honoring
+    ``warmup_steps``. During warmup, the temperature stays at ``start``.
+
+    Parameters
+    - start: Temperature at the beginning (``p = 0``). Must be ``> 0``.
+    - end: Temperature at the end (``p = 1``). Must be ``> 0``.
+    - total_steps: Default total steps for the schedule. Can be overridden per
+      call using the ``total=`` kwarg.
+    - warmup_steps: Number of initial steps to hold at ``start`` before
+      annealing begins. Must be ``>= 0``.
+    """
+
+    def __init__(
+        self,
+        start: float,
+        end: float,
+        *,
+        total_steps: int | None = None,
+        warmup_steps: int = 0,
+    ) -> None:
+        s = float(start)
+        e = float(end)
+        if not math.isfinite(s) or s <= 0.0:
+            raise ValueError("start must be a positive finite float")
+        if not math.isfinite(e) or e <= 0.0:
+            raise ValueError("end must be a positive finite float")
+        if total_steps is not None and int(total_steps) <= 0:
+            raise ValueError("total_steps must be >= 1 when provided")
+        if int(warmup_steps) < 0:
+            raise ValueError("warmup_steps must be >= 0")
+
+        self.start: float = s
+        self.end: float = e
+        self.total_steps: int | None = int(total_steps) if total_steps is not None else None
+        self.warmup_steps: int = int(warmup_steps)
+
+    def __repr__(self) -> str:  # pragma: no cover - cosmetic
+        return (
+            "CosineTemperature("
+            f"start={self.start}, end={self.end}, "
+            f"total_steps={self.total_steps}, warmup_steps={self.warmup_steps})"
+        )
+
+    def _progress_from_step(self, step: int, total: int) -> float:
+        # Compute normalized progress in [0, 1] after warmup.
+        s = max(0, int(step))
+        t = max(1, int(total))
+        w = max(0, self.warmup_steps)
+        if s <= w:
+            return 0.0
+        # Remaining steps after warmup; guard against division by zero.
+        remain = max(1, t - w)
+        p = (s - w) / float(remain)
+        # Clamp to [0, 1] to be robust against overrun.
+        if p < 0.0:
+            return 0.0
+        if p > 1.0:
+            return 1.0
+        return p
+
+    @staticmethod
+    def _cosine_mix(start: float, end: float, p01: float) -> float:
+        # Clamp progress defensively.
+        p = 0.0 if p01 < 0.0 else (1.0 if p01 > 1.0 else p01)
+        return end + (start - end) * 0.5 * (1.0 + math.cos(math.pi * p))
+
+    def __call__(
+        self,
+        step: int | None = None,
+        *,
+        total: int | None = None,
+        progress: float | None = None,
+    ) -> float:
+        if progress is not None:
+            return self._cosine_mix(self.start, self.end, float(progress))
+
+        # Derive progress from steps; prefer call-time total over default.
+        resolved_total = int(total) if total is not None else self.total_steps
+        if resolved_total is None or step is None:
+            raise ValueError(
+                "CosineTemperature requires either progress in [0,1] or both step and total"
+            )
+
+        p = self._progress_from_step(int(step), int(resolved_total))
+        # During warmup (_progress_from_step returns 0), hold at start when step <= warmup
+        if int(step) <= self.warmup_steps:
+            return float(self.start)
+        return self._cosine_mix(self.start, self.end, p)
+
+    def as_tensor(
+        self,
+        step: int | None = None,
+        *,
+        total: int | None = None,
+        progress: float | None = None,
+        device: DeviceT | None = None,
+        dtype: DTypeT | None = None,
+    ) -> TensorT:
+        val = self(step, total=total, progress=progress)
+        torch = _require_torch()
+        return cast("TensorT", torch.tensor(val, device=device, dtype=dtype))

nervecode/core/trace.py

@@ -0,0 +1,166 @@
+"""Trace data structures for coding.
+
+This module defines lightweight, immutable containers that carry the rich
+per-layer record produced during coding. Traces support arbitrary leading
+dimensions so they can represent batch-only, token-like, or pooled layouts.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from .types import SoftCode
+
+__all__ = ["CodingTrace"]
+
+
+@dataclass(frozen=True)
+class CodingTrace:
+    """Per-layer coding trace.
+
+    Carries the reduced activation, reduction metadata, nearest-center and
+    commitment distances, chosen indices, and the corresponding ``SoftCode``.
+
+    Fields
+    - ``reduced`` (Tensor): reduced activation with shape ``(..., D)``.
+    - ``reduction_meta`` (dict): metadata describing the applied reduction
+      (e.g., method, parameters). May be empty when no reduction is applied.
+    - ``nearest_center_distances`` (Tensor): distance to the selected center
+      per position with shape matching the leading dimensions ``...``.
+    - ``chosen_center_indices`` (Tensor): index of the selected/nearest center
+      per position with shape matching the leading dimensions ``...``.
+    - ``commitment_distances`` (Tensor): commitment distance per position with
+      shape matching the leading dimensions ``...``.
+    - ``soft_code`` (SoftCode): soft assignment statistics over ``K`` codes
+      with probabilities of shape ``(..., K)``.
+
+    All scalar-like fields must match the leading shape of ``soft_code.probs``
+    exactly (i.e., ``tensor.shape == soft_code.probs.shape[:-1]``).
+    """
+
+    reduced: Any
+    reduction_meta: dict[str, Any]
+    nearest_center_distances: Any
+    chosen_center_indices: Any
+    commitment_distances: Any
+    soft_code: SoftCode
+
+    def __post_init__(self) -> None:
+        # Import torch lazily to keep import-time dependencies lightweight.
+        from typing import Any
+        from typing import cast as _cast
+
+        try:
+            import torch
+        except Exception:  # pragma: no cover - torch may be absent in some envs
+            torch = _cast(Any, None)
+
+        # Basic validation of required components
+        if not isinstance(self.soft_code, SoftCode):  # pragma: no cover - defensive
+            raise TypeError("soft_code must be a SoftCode instance")
+
+        # Validate reduced tensor shape
+        if torch is not None and not isinstance(
+            self.reduced, torch.Tensor
+        ):  # pragma: no cover - defensive
+            raise TypeError("reduced must be a torch.Tensor")
+        if getattr(self.reduced, "ndim", None) is None or self.reduced.ndim < 1:
+            raise ValueError("reduced must have at least one dimension (..., D)")
+        if int(self.reduced.shape[-1]) <= 0:
+            raise ValueError("final reduced dimension D must be >= 1")
+
+        # If torch isn't available, skip dtype/shape validations that require it.
+        if torch is None:  # pragma: no cover
+            return
+
+        lead_shape = tuple(int(s) for s in self.soft_code.leading_shape)
+        reduced_lead = tuple(int(s) for s in self.reduced.shape[:-1])
+        if reduced_lead != lead_shape:
+            raise ValueError(
+                f"reduced.leading_shape {reduced_lead} must match "
+                f"soft_code.leading_shape {lead_shape}"
+            )
+
+        def _check_scalar_tensor(name: str, value: Any, *, require_integer: bool = False) -> None:
+            if not isinstance(value, torch.Tensor):
+                raise TypeError(f"{name} must be a torch.Tensor")
+            if tuple(int(s) for s in value.shape) != lead_shape:
+                raise ValueError(
+                    f"{name} must have shape {lead_shape} to match soft_code.leading_shape"
+                )
+            if require_integer:
+                if torch.is_floating_point(value) or value.dtype == torch.bool:
+                    raise TypeError(f"{name} must use an integer dtype")
+            else:
+                if value.dtype == torch.bool:
+                    raise TypeError(f"{name} must be numeric, not bool")
+
+        _check_scalar_tensor("nearest_center_distances", self.nearest_center_distances)
+        _check_scalar_tensor(
+            "chosen_center_indices", self.chosen_center_indices, require_integer=True
+        )
+        _check_scalar_tensor("commitment_distances", self.commitment_distances)
+
+    @property
+    def leading_shape(self) -> tuple[int, ...]:
+        """Leading shape shared by all per-position fields."""
+
+        return tuple(int(s) for s in self.soft_code.leading_shape)
+
+    @property
+    def reduced_dim(self) -> int:
+        """Size of the final reduced dimension ``D``."""
+
+        return int(self.reduced.shape[-1])
+
+    # --- Sample-level reduction views -------------------------------------------
+    def sample_reduced_surprise(self, *, reduction: str = "mean") -> Any:
+        """Return a per-sample surprise tensor by reducing extra leading dims.
+
+        Many wrappers operate over activations with multiple leading dimensions
+        (e.g., batch and sequence). Aggregation across layers requires a common
+        per-sample view. This helper reduces the per-position surprise stored on
+        the associated ``SoftCode`` to a vector over samples by collapsing all
+        leading dimensions after the first (typically sequence/time) using the
+        specified reduction.
+
+        Parameters
+        - reduction: Current supported value is ``"mean"``. Additional
+          reductions may be added later.
+
+        Returns
+        - A tensor with shape ``(B,)`` where ``B`` is the size of the first
+          leading dimension. In environments without torch or when an
+          interpretable surprise is unavailable, this may return ``None`` at
+          runtime; the return type is ``Any`` to remain flexible during import
+          without a hard torch dependency.
+        """
+
+        # Import torch lazily to avoid hard dependency at import time.
+        try:
+            import torch
+        except Exception:  # pragma: no cover - environments without torch
+            return None
+
+        # Prefer a combined surprise when present; otherwise use best_length.
+        cand = self.soft_code.combined_surprise
+        if cand is None:
+            cand = self.soft_code.best_length
+        if not isinstance(cand, torch.Tensor):
+            return None
+
+        # Scalar-like can't be mapped to samples reliably; skip.
+        if cand.ndim == 0:  # pragma: no cover - uncommon layout
+            return None
+
+        if cand.ndim == 1:
+            return cand
+
+        # Reduce all leading dims after the first to obtain a per-sample vector.
+        if reduction == "mean":
+            dims = tuple(range(1, int(cand.ndim)))
+            return cand.mean(dim=dims)
+        # Fallback: default to mean for unknown identifiers to remain fail-open.
+        dims = tuple(range(1, int(cand.ndim)))
+        return cand.mean(dim=dims)