nervecode 0.1.0__py3-none-any.whl

nervecode/scoring/aggregator.py

@@ -0,0 +1,369 @@
+"""Surprise aggregation utilities.
+
+MVP provides simple, explicit aggregation strategies across wrapped layers:
+- mean, max, and fixed weighted combinations of per-layer surprise signals.
+
+The primary entry point, ``mean_surprise(...)``, accepts a collection of
+per-layer traces or soft-code objects and returns a per-sample aggregated
+surprise tensor. Inputs are interpreted best-effort to preserve fail-open
+behavior:
+
+- ``CodingTrace``: use ``soft_code.combined_surprise`` if available, otherwise
+  fall back to ``soft_code.best_length``.
+- ``SoftCode``: use ``combined_surprise`` if available, otherwise ``best_length``.
+- ``torch.Tensor``: treated as an already-computed per-sample surprise signal.
+
+Entries without an interpretable per-sample signal are skipped. When no valid
+signals remain, ``None`` is returned. When multiple layers are provided,
+signals are first reduced to a sample-level view so that wrappers with
+different leading dimensions (e.g., ``(B, T)`` vs. ``(B,)``) can be combined.
+Layers that cannot provide a sample-level view are skipped.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from typing import Any, cast
+
+try:  # Keep import-time behavior tolerant in environments without torch
+    import torch
+except Exception:  # pragma: no cover - torch is a project dependency in tests
+    torch = cast(Any, None)
+
+try:
+    from nervecode.core import CodingTrace, SoftCode  # re-exported types
+except Exception:  # pragma: no cover - available during normal package use
+    CodingTrace = object  # type: ignore[misc,assignment]
+    SoftCode = object  # type: ignore[misc,assignment]
+
+from .types import AggregatedSurprise
+
+__all__ = ["AggregatedSurprise", "max_surprise", "mean_surprise", "weighted_surprise"]
+
+
+def mean_surprise(
+    traces: Mapping[str, Any] | Iterable[Any],
+) -> AggregatedSurprise | None:
+    """Return the mean-aggregated surprise across layers.
+
+    Parameters
+    - traces: Mapping from layer names to trace-like objects or an iterable of
+      such entries. Each entry may be a ``CodingTrace``, ``SoftCode``, or a
+      per-sample ``torch.Tensor``.
+
+    Returns
+    - ``AggregatedSurprise`` with the aggregated per-sample signal and basic
+      metadata, or ``None`` when no valid per-sample signals are available.
+    """
+
+    if torch is None:  # pragma: no cover - defensive
+        return None
+
+    # Normalize input to an iterable of values
+    values: Iterable[Any] = traces.values() if isinstance(traces, Mapping) else traces
+
+    signals: list[torch.Tensor] = []
+    ref_shape: tuple[int, ...] | None = None
+
+    for obj in values:
+        t = _as_sample_surprise_tensor(obj)
+        if t is None:
+            continue
+        # Use the shape of the first valid tensor as a reference
+        shape = tuple(int(s) for s in t.shape)
+        if ref_shape is None:
+            ref_shape = shape
+            signals.append(t)
+        else:
+            # Only aggregate signals with the same per-sample shape.
+            if shape != ref_shape:
+                continue
+            signals.append(t)
+
+    if not signals:
+        return None
+
+    # Bring all tensors to the device/dtype of the first entry for stacking
+    first = signals[0]
+    device = getattr(first, "device", None)
+    dtype = first.dtype if hasattr(first, "dtype") else None
+    aligned: list[torch.Tensor] = []
+    for s in signals:
+        s2 = s
+        if dtype is not None and getattr(s2, "dtype", None) is not dtype:
+            s2 = s2.to(dtype=dtype)
+        if device is not None and getattr(s2, "device", None) != device:
+            s2 = s2.to(device=device)
+        aligned.append(s2)
+
+    stacked = torch.stack(aligned, dim=0)
+    agg = stacked.mean(dim=0)
+    return AggregatedSurprise(
+        surprise=agg,
+        method="mean",
+        num_layers=len(aligned),
+        details=None,
+    )
+
+
+def max_surprise(
+    traces: Mapping[str, Any] | Iterable[Any],
+) -> AggregatedSurprise | None:
+    """Return the max-aggregated surprise across layers.
+
+    Mirrors ``mean_surprise`` for input handling and sample-level reduction but
+    combines participating layer signals using a max across layers instead of a
+    mean. Returns ``None`` when no valid per-sample signals are available.
+    """
+
+    if torch is None:  # pragma: no cover - defensive
+        return None
+
+    values: Iterable[Any] = traces.values() if isinstance(traces, Mapping) else traces
+
+    signals: list[torch.Tensor] = []
+    ref_shape: tuple[int, ...] | None = None
+
+    for obj in values:
+        t = _as_sample_surprise_tensor(obj)
+        if t is None:
+            continue
+        shape = tuple(int(s) for s in t.shape)
+        if ref_shape is None:
+            ref_shape = shape
+            signals.append(t)
+        else:
+            if shape != ref_shape:
+                continue
+            signals.append(t)
+
+    if not signals:
+        return None
+
+    first = signals[0]
+    device = getattr(first, "device", None)
+    dtype = first.dtype if hasattr(first, "dtype") else None
+    aligned: list[torch.Tensor] = []
+    for s in signals:
+        s2 = s
+        if dtype is not None and getattr(s2, "dtype", None) is not dtype:
+            s2 = s2.to(dtype=dtype)
+        if device is not None and getattr(s2, "device", None) != device:
+            s2 = s2.to(device=device)
+        aligned.append(s2)
+
+    stacked = torch.stack(aligned, dim=0)
+    agg = torch.max(stacked, dim=0).values
+    return AggregatedSurprise(
+        surprise=agg,
+        method="max",
+        num_layers=len(aligned),
+        details=None,
+    )
+
+
+def weighted_surprise(
+    traces: Mapping[str, Any] | Iterable[Any],
+    *,
+    weights: Mapping[str, float] | Iterable[float] | None = None,
+    normalize: bool = True,
+) -> AggregatedSurprise | None:
+    """Return a fixed weighted aggregation across layers.
+
+    The function mirrors ``mean_surprise`` for input handling and sample-level
+    reduction, but combines participating layer signals using explicit fixed
+    weights instead of a uniform mean. This function does not learn weights.
+
+    Parameters
+    - traces: Mapping from layer names to trace-like objects or an iterable of
+      such entries. Each entry may be a ``CodingTrace``, ``SoftCode``, or a
+      per-sample ``torch.Tensor``.
+    - weights: Either a mapping from layer name to weight (preferred when
+      ``traces`` is a mapping) or an iterable of weights aligned with the order
+      of ``traces``. When omitted, all included layers use weight ``1.0``.
+    - normalize: When ``True`` (default), divides by the sum of the included
+      weights so that outputs are comparable to a mean when weights are equal.
+
+    Returns
+    - ``AggregatedSurprise`` with aggregation metadata including the effective
+      weights used, or ``None`` when no valid per-sample signals are available
+      or the total weight after filtering is zero.
+    """
+
+    if torch is None:  # pragma: no cover - defensive
+        return None
+
+    # Normalize to lists of (name, value) to preserve a stable order for
+    # iterable inputs and enable name-based weight lookup when available.
+    items: list[tuple[str | None, Any]]
+    if isinstance(traces, Mapping):
+        items = [(str(k), v) for k, v in traces.items()]
+    else:
+        items = [(None, v) for v in traces]
+
+    # Prepare a parallel list of sample-level tensors alongside resolved weights.
+    signals: list[torch.Tensor] = []
+    resolved_weights: list[float] = []
+    names: list[str | None] = []
+    ref_shape: tuple[int, ...] | None = None
+
+    # Helper to get a weight for an item given its index and (optional) name.
+    weights_map: Mapping[str, float] | None
+    weights_seq: list[float] | None
+    if weights is None:
+        weights_map = None
+        weights_seq = None
+    elif isinstance(weights, Mapping):
+        # Use a plain mapping lookup for name-keyed inputs
+        weights_map = {str(k): float(v) for k, v in weights.items()}
+        weights_seq = None
+    else:
+        # Materialize iterable weights to support index-aligned access
+        try:
+            weights_seq = [float(v) for v in list(weights)]
+        except Exception:
+            weights_seq = None
+        weights_map = None
+
+    def _weight_for(idx: int, name: str | None) -> float:
+        if weights_map is not None and name is not None:
+            return float(weights_map.get(name, 1.0))
+        if weights_seq is not None and 0 <= idx < len(weights_seq):
+            return float(weights_seq[idx])
+        return 1.0
+
+    for idx, (name, obj) in enumerate(items):
+        t = _as_sample_surprise_tensor(obj)
+        if t is None:
+            continue
+        shape = tuple(int(s) for s in t.shape)
+        if ref_shape is None:
+            ref_shape = shape
+        if shape != ref_shape:
+            # Skip mismatched sample shapes to preserve fail-open behavior.
+            continue
+        w_i = _weight_for(idx, name)
+        # Skip strictly non-positive weights to avoid divide-by-zero surprises
+        # and preserve user intent (e.g., selectively disabling layers).
+        if not (w_i > 0.0):
+            continue
+        names.append(name)
+        signals.append(t)
+        resolved_weights.append(float(w_i))
+
+    if not signals:
+        return None
+
+    # Align device/dtype across participating tensors.
+    first = signals[0]
+    device = getattr(first, "device", None)
+    dtype = first.dtype if hasattr(first, "dtype") else None
+    aligned: list[torch.Tensor] = []
+    for s in signals:
+        s2 = s
+        if dtype is not None and getattr(s2, "dtype", None) is not dtype:
+            s2 = s2.to(dtype=dtype)
+        if device is not None and getattr(s2, "device", None) != device:
+            s2 = s2.to(device=device)
+        aligned.append(s2)
+
+    # Convert weights to a tensor on the target device/dtype for safe math.
+    w_vec = torch.tensor(resolved_weights, dtype=dtype or torch.float32, device=device)
+    total = float(w_vec.sum().item())
+    if total <= 0.0:
+        return None
+
+    stacked = torch.stack(aligned, dim=0)
+    # Weighted sum along layer dimension
+    # Shape: (layers, ...) -> (...) via tensordot/broadcasted mul then sum.
+    weighted = stacked * w_vec.view(-1, *([1] * (stacked.ndim - 1)))
+    agg = weighted.sum(dim=0)
+    if normalize:
+        agg = agg / total
+
+    # Provide simple details: per-layer names and weights actually used.
+    used: list[tuple[str | None, float]] = list(zip(names, resolved_weights))
+    details: dict[str, Any] = {
+        "weights": used,
+        "normalized": bool(normalize),
+        "sum_weights": total,
+    }
+
+    return AggregatedSurprise(
+        surprise=agg,
+        method="weighted",
+        num_layers=len(aligned),
+        details=details,
+    )
+
+
+def _as_sample_surprise_tensor(obj: Any) -> torch.Tensor | None:
+    """Return a per-sample surprise tensor for a supported input object.
+
+    All supported inputs are reduced to a sample-level vector by collapsing
+    leading dimensions after the first (using mean) when necessary. This allows
+    combining wrappers that produce position- or token-level signals with
+    batch-only signals.
+
+    Supported inputs:
+    - CodingTrace: use ``sample_reduced_surprise()`` when available; otherwise
+      prefer ``soft_code.combined_surprise``, else ``soft_code.best_length``.
+    - SoftCode: prefer combined_surprise, else best_length; reduce to ``(B,)``
+      when rank > 1.
+    - torch.Tensor: reduce to ``(B,)`` when rank > 1.
+    """
+
+    if torch is None:  # pragma: no cover - defensive
+        return None
+
+    # Torch tensor path
+    if hasattr(torch, "Tensor") and isinstance(obj, torch.Tensor):
+        return _reduce_to_sample(obj)
+
+    # SoftCode path
+    if isinstance(obj, SoftCode):
+        cand = getattr(obj, "combined_surprise", None)
+        if cand is None:
+            cand = getattr(obj, "best_length", None)
+        if isinstance(cand, torch.Tensor):
+            return _reduce_to_sample(cand)
+        return None
+
+    # CodingTrace path
+    if isinstance(obj, CodingTrace):
+        # Prefer a sample-level view provided by the trace itself.
+        view = getattr(obj, "sample_reduced_surprise", None)
+        if callable(view):
+            try:
+                out = view()
+                if isinstance(out, torch.Tensor):
+                    return out
+            except Exception:
+                pass
+        sc = getattr(obj, "soft_code", None)
+        if isinstance(sc, SoftCode):
+            return _as_sample_surprise_tensor(sc)
+        return None
+
+    # Mapping path: allow passing a dict-like with a direct 'surprise' tensor
+    if isinstance(obj, Mapping):
+        # Common keys we might accept in future; keep minimal for MVP.
+        cand = obj.get("surprise") if hasattr(obj, "get") else None
+        if isinstance(cand, torch.Tensor):
+            return _reduce_to_sample(cand)
+
+    return None
+
+
+def _reduce_to_sample(t: torch.Tensor) -> torch.Tensor:
+    """Reduce a tensor to a per-sample vector by averaging extra leading dims.
+
+    - ``(B,)`` -> returned unchanged
+    - ``(B, T, ...)`` -> mean over dims 1..N-1
+    - ``()`` (scalar) -> returned as-is (caller may choose to skip)
+    """
+
+    if getattr(t, "ndim", 0) <= 1:
+        return t
+    dims = tuple(range(1, int(t.ndim)))
+    return t.mean(dim=dims)