nervecode 0.1.0__py3-none-any.whl

nervecode/training/loss.py

@@ -0,0 +1,188 @@
+"""Auxiliary coding loss that consumes CodingTrace objects.
+
+This module implements a lightweight training-time loss that reads existing
+``CodingTrace`` objects produced by wrapped layers. It avoids recomputing
+distances or other intermediates from raw activations by relying on the
+per-position surprise already stored on the associated ``SoftCode`` inside each
+trace. When multiple layers are provided, their per-sample surprise signals are
+aggregated using the default mean strategy from ``nervecode.scoring`` and then
+reduced across the batch to yield a scalar loss.
+
+The loss intentionally remains minimal in the MVP: it performs a mean reduction
+over samples. Follow-up tasks extend it with configurable term weights and
+clear per-layer breakdowns.
+"""
+
+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
+    from torch import nn
+except Exception:  # pragma: no cover - torch is a project dependency in tests
+    torch = cast(Any, None)
+    nn = 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 nervecode.scoring import mean_surprise
+
+__all__ = ["CodingLoss"]
+
+
+class _StubModule:  # pragma: no cover - used only when torch is unavailable
+    pass
+
+
+_Base = nn.Module if nn is not None else _StubModule
+
+
+class CodingLoss(_Base):  # type: ignore[misc,valid-type]
+    """Compute a scalar coding loss from traces without recomputation.
+
+    Inputs may be a single ``CodingTrace`` (or ``SoftCode``/tensor), an
+    iterable of such objects, or a mapping from layer names to objects.
+
+    The loss has two conceptual parts:
+    - an assignment/sharpness term derived from the per-position surprise
+      on ``SoftCode`` (aggregated via ``mean_surprise``), and
+    - an optional commitment term that uses the nearest-center distance
+      already stored on each ``CodingTrace``.
+
+    Both terms are reduced to a per-sample view and averaged across layers;
+    the final scalar is a weighted sum with ``assign_weight`` and
+    ``commit_weight``. By default the commitment term weight is zero to keep
+    backward compatibility with earlier behavior.
+    """
+
+    def __init__(self, *, assign_weight: float = 1.0, commit_weight: float = 0.0) -> None:
+        if torch is None:  # pragma: no cover - defensive
+            raise RuntimeError("CodingLoss requires PyTorch to be installed")
+        super().__init__()
+        self.assign_weight: float = float(assign_weight)
+        self.commit_weight: float = float(commit_weight)
+
+    def forward(self, traces: Any) -> torch.Tensor:
+        """Return weighted mean-aggregated coding loss across samples.
+
+        Parameters
+        - traces: A ``CodingTrace`` object (or ``SoftCode``/tensor), an
+          iterable of such objects, or a mapping from layer names to objects.
+
+        Returns
+        - A scalar ``torch.Tensor`` equal to the weighted sum of the per-sample
+          aggregated assignment term and, when enabled, the aggregated
+          commitment term.
+        """
+
+        if torch is None:  # pragma: no cover - defensive
+            raise RuntimeError("CodingLoss requires PyTorch to be installed")
+
+        # Normalize input to what mean_surprise and our commitment aggregator expect
+        payload: Iterable[Any] | Mapping[str, Any]
+        if isinstance(traces, Mapping) or (
+            isinstance(traces, Iterable)
+            and not isinstance(traces, (torch.Tensor, SoftCode, CodingTrace))
+        ):
+            payload = traces
+        else:
+            # Single object: wrap into a list for aggregation
+            payload = [traces]
+
+        # Assignment/sharpness component via mean_surprise (SoftCode-based)
+        agg = mean_surprise(payload)
+        if agg is None or not hasattr(agg, "surprise"):
+            raise ValueError("No valid surprise signals found in provided traces")
+        s = cast(torch.Tensor, agg.surprise)
+        assign_term = s.mean()
+
+        # Commitment component via nearest-center distances stored on traces
+        commit_vec: torch.Tensor | None = _mean_commitment_across_layers(payload)
+
+        if self.commit_weight != 0.0 and commit_vec is not None:
+            commit_mean = commit_vec.mean()
+            return (
+                assign_term.to(dtype=commit_mean.dtype, device=commit_mean.device)
+                * self.assign_weight
+                + commit_mean * self.commit_weight
+            )
+        # Default: only assignment term
+        return assign_term * self.assign_weight
+
+
+def _mean_commitment_across_layers(
+    traces: Iterable[Any] | Mapping[str, Any],
+) -> torch.Tensor | None:
+    """Return per-sample mean-aggregated commitment distance across layers.
+
+    This helper scans the provided collection for ``CodingTrace`` objects,
+    extracts their per-position ``commitment_distances``, reduces them to a
+    per-sample vector by averaging extra leading dimensions, and averages
+    across layers. When no layer provides commitment distances, returns
+    ``None``.
+    """
+
+    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:
+        # Only CodingTrace carries commitment distances in the MVP
+        if not isinstance(obj, CodingTrace):
+            continue
+        d = getattr(obj, "commitment_distances", None)
+        if not isinstance(d, torch.Tensor):
+            continue
+        t = _reduce_to_sample(d)
+        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:
+                # Skip mismatched sample shapes to preserve fail-open behavior
+                continue
+            signals.append(t)
+
+    if not signals:
+        return None
+
+    # Align dtype/device to the first signal for stable aggregation
+    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)
+    # Mean over layers yields a per-sample vector
+    return stacked.mean(dim=0)
+
+
+def _reduce_to_sample(t: torch.Tensor) -> torch.Tensor:
+    """Reduce a tensor to a per-sample vector by averaging extra leading dims.
+
+    Mirrors the reduction used by the scoring aggregator so that both
+    assignment and commitment terms are combined consistently across layers.
+    """
+
+    if getattr(t, "ndim", 0) <= 1:
+        return t
+    dims = tuple(range(1, int(t.ndim)))
+    return t.mean(dim=dims)

nervecode/training/updaters.py

@@ -0,0 +1,168 @@
+"""Experimental codebook update helpers (EMA / hybrid).
+
+These helpers are optional, off-by-default training utilities that update
+codebook centers using exponential moving averages driven by the latest
+coding trace. They are provided for exploration only — the default and
+recommended baseline remains standard gradient-based optimization via
+``CodingLoss`` and your optimizer.
+
+Design goals
+- Observe-only wrappers and existing training code remain unchanged unless an
+  updater is explicitly instantiated and called by the user.
+- Safe numerics: clamp divisors, support empty usages gracefully.
+- Torch-first: implemented as small ``nn.Module``s with buffers so they move
+  with devices and can be checkpointed when needed.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal, cast
+
+try:  # Keep import-time failure tolerant when torch isn't installed
+    import torch
+    from torch import nn
+except Exception:  # pragma: no cover - torch is a project dependency in tests
+    torch = cast(Any, None)
+    nn = cast(Any, None)
+
+try:
+    from nervecode.core import CodingTrace
+    from nervecode.core.codebook import Codebook
+except Exception:  # pragma: no cover - available during normal package use
+    CodingTrace = object  # type: ignore[misc,assignment]
+    Codebook = object  # type: ignore[misc,assignment]
+
+__all__ = ["EmaCodebookUpdater"]
+
+
+class _StubModule:  # pragma: no cover - used only when torch is unavailable
+    pass
+
+
+_Base = nn.Module if nn is not None else _StubModule
+
+
+class EmaCodebookUpdater(_Base):  # type: ignore[misc,valid-type]
+    """Update codebook centers using an EMA of batch-weighted means.
+
+    This updater maintains per-code EMA counts and summed activations and sets
+    centers to ``ema_sum / (ema_count + eps)`` after each update call. The batch
+    statistics are computed from a ``CodingTrace`` using either soft
+    probabilities or hard assignments.
+
+    Parameters
+    - codebook: The ``Codebook`` whose centers are updated in-place.
+    - decay: Exponential decay factor in ``[0, 1)``. Higher means slower update
+      toward the batch statistics. ``decay=0`` uses the current batch only.
+    - assignment: ``"soft"`` (default) to use soft probabilities, or ``"hard"``
+      to use one-hot assignments based on argmax indices.
+    - eps: Small positive constant to avoid division by zero.
+
+    Notes
+    - Updates run under ``torch.no_grad()`` and preserve ``requires_grad`` on
+      ``codebook.centers`` so gradient-based training can continue to operate
+      (hybrid usage).
+    - The updater holds its own buffers (ema_count, ema_sum) sized from the
+      attached codebook at construction time. If the codebook is reinitialized
+      with a different shape, construct a new updater.
+    """
+
+    def __init__(
+        self,
+        codebook: Codebook,
+        *,
+        decay: float = 0.99,
+        assignment: Literal["soft", "hard"] = "soft",
+        eps: float = 1e-12,
+    ) -> None:
+        if torch is None:  # pragma: no cover - defensive
+            raise RuntimeError("EmaCodebookUpdater requires PyTorch to be installed")
+        super().__init__()
+
+        if decay < 0.0 or decay >= 1.0 or not float(decay) == decay:
+            raise ValueError("decay must be a float in [0, 1)")
+        if assignment not in ("soft", "hard"):
+            raise ValueError("assignment must be 'soft' or 'hard'")
+        if eps <= 0.0:
+            raise ValueError("eps must be positive")
+
+        self.codebook = codebook
+        self.decay: float = float(decay)
+        self.assignment: Literal["soft", "hard"] = assignment
+        self.eps: float = float(eps)
+
+        K = int(codebook.K)
+        D = int(codebook.code_dim)
+        # EMA state mirrors codebook shape: counts per code and sums per code x dim
+        self._ema_count: torch.Tensor
+        self._ema_sum: torch.Tensor
+        self.register_buffer("_ema_count", torch.zeros(K), persistent=True)
+        self.register_buffer("_ema_sum", torch.zeros(K, D), persistent=True)
+
+    def reset_state(self) -> None:
+        """Reset EMA state to zeros (no prior influence)."""
+        if torch is None:  # pragma: no cover - defensive
+            return
+        with torch.no_grad():
+            self._ema_count.zero_()
+            self._ema_sum.zero_()
+
+    def update_from_trace(self, trace: CodingTrace) -> None:
+        """Update EMA statistics and set codebook centers for a batch.
+
+        Expects a ``CodingTrace`` produced by the same codebook. The reduced
+        activations and the ``SoftCode`` probabilities/indices are used to
+        compute per-code batch counts and sums.
+        """
+
+        # Validate shape alignment in a torch-first way
+        probs = trace.soft_code.probs
+        x = trace.reduced
+
+        if not isinstance(probs, torch.Tensor) or not isinstance(x, torch.Tensor):
+            raise TypeError("trace must carry torch.Tensor probabilities and reduced activations")
+        if probs.ndim < 1 or x.ndim < 1:
+            raise ValueError("trace tensors must have at least one dimension")
+        if int(probs.shape[-1]) != int(self.codebook.K):
+            raise ValueError("trace code dimension K must match attached codebook")
+        if int(x.shape[-1]) != int(self.codebook.code_dim):
+            raise ValueError("trace reduced dim D must match attached codebook")
+
+        with torch.no_grad():
+            # Flatten leading dims to (N, ...)
+            N = int(x.numel() // x.shape[-1])
+            x_flat = x.reshape(N, -1)  # (N, D)
+            p_flat = probs.reshape(N, -1)  # (N, K)
+
+            if self.assignment == "hard":
+                # One-hot from argmax indices
+                idx = torch.argmax(p_flat, dim=-1)  # (N,)
+                w = torch.zeros_like(p_flat)
+                w.scatter_(1, idx.unsqueeze(-1), 1.0)
+            else:
+                # Soft probabilities (already sum to 1 across codes)
+                w = p_flat
+
+            # Per-code batch counts and sums
+            batch_count = w.sum(dim=0)  # (K,)
+            batch_sum = w.t() @ x_flat  # (K, D)
+
+            # Align EMA state to the same device/dtype as incoming batch
+            device = x_flat.device
+            dtype = x_flat.dtype
+            if self._ema_count.device != device or self._ema_count.dtype != dtype:
+                self._ema_count = self._ema_count.to(device=device, dtype=dtype)
+            if self._ema_sum.device != device or self._ema_sum.dtype != dtype:
+                self._ema_sum = self._ema_sum.to(device=device, dtype=dtype)
+
+            d = self.decay
+            one_minus = 1.0 - d
+            self._ema_count.mul_(d).add_(batch_count.to(dtype=dtype), alpha=one_minus)
+            self._ema_sum.mul_(d).add_(batch_sum.to(dtype=dtype), alpha=one_minus)
+
+            # Update centers: ema_sum / (ema_count + eps)
+            denom = (self._ema_count + self.eps).unsqueeze(-1)  # (K, 1)
+            new_centers = self._ema_sum / denom
+
+            # In-place update preserves requires_grad on the parameter
+            self.codebook.centers.detach().copy_(new_centers)

nervecode/utils/__init__.py

@@ -0,0 +1,14 @@
+"""Small utilities shared across the package.
+
+Utility helpers should remain lightweight and free of heavy dependencies.
+"""
+
+from __future__ import annotations
+
+from .seed import seed_everything, seed_from_env, temp_seed
+
+__all__: list[str] = [
+    "seed_everything",
+    "seed_from_env",
+    "temp_seed",
+]

nervecode/utils/overhead.py

@@ -0,0 +1,177 @@
+"""Overhead estimators for pooled Conv2d coding.
+
+This module provides lightweight, deterministic estimators for the compute and
+memory overhead of the pooled-convolutional coding path used by
+``nervecode.layers.conv.CodingConv2d``. The estimates are torch-agnostic and
+intended for documentation, sanity checks, and unit tests — not for replacing
+runtime profiling.
+
+Conventions
+- We report multiply-accumulate operations (MACs) rather than FLOPs. A single
+  multiply followed by an add counts as one MAC. The baseline Conv2d MACs use
+  the common approximation: ``B * H_out * W_out * C_out * (C_in * kH * kW)``.
+- Coding overhead counts are broken down into: global-average pooling over
+  spatial dimensions, optional linear projection (when ``coding_dim < C_out``),
+  and codebook distance + soft-assignment operations of size ``(B, D, K)``.
+- Memory estimates cover transient activations in the coding path (e.g., the
+  pooled vector ``(B, D)`` and soft-code ``(B, K)``) and parameter overhead
+  introduced by the wrapper (codebook centers and optional projection matrix).
+
+These estimates aim to be simple and stable across environments. Real kernels
+and library implementations may differ by constant factors, fusions, and cache
+effects; treat the results as order-of-magnitude guidance.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class PooledConvConfig:
+    """Configuration for a 2D convolution with pooled coding.
+
+    All dimensions are positive integers. ``kernel_size``, ``stride``,
+    ``padding``, and ``dilation`` accept symmetric pairs as ``(h, w)``.
+
+    Fields
+    - batch: batch size ``B``
+    - in_channels: input channels ``C_in``
+    - out_channels: output channels ``C_out``
+    - in_size: input spatial size ``(H_in, W_in)``
+    - kernel_size: kernel size ``(kH, kW)``
+    - stride: stride ``(sH, sW)``
+    - padding: padding ``(pH, pW)``
+    - dilation: dilation ``(dH, dW)``
+    - coding_dim: pooled coding dimension ``D``; defaults to ``C_out``
+    - K: codebook size
+    - dtype_bytes: bytes per float element for activation estimates (default: 4)
+    """
+
+    batch: int
+    in_channels: int
+    out_channels: int
+    in_size: tuple[int, int]
+    kernel_size: tuple[int, int] = (3, 3)
+    stride: tuple[int, int] = (1, 1)
+    padding: tuple[int, int] = (0, 0)
+    dilation: tuple[int, int] = (1, 1)
+    coding_dim: int | None = None
+    K: int = 16
+    dtype_bytes: int = 4
+
+    def output_size(self) -> tuple[int, int]:
+        """Compute the output spatial size (H_out, W_out) for the conv2d.
+
+        Uses the standard convolution formula with floor division.
+        """
+
+        H_in, W_in = self.in_size
+        kH, kW = self.kernel_size
+        sH, sW = self.stride
+        pH, pW = self.padding
+        dH, dW = self.dilation
+
+        H_out = (H_in + 2 * pH - dH * (kH - 1) - 1) // sH + 1
+        W_out = (W_in + 2 * pW - dW * (kW - 1) - 1) // sW + 1
+        return int(H_out), int(W_out)
+
+    def resolved_coding_dim(self) -> int:
+        """Return effective coding dimension ``D`` (defaults to ``C_out``)."""
+
+        return int(self.out_channels if self.coding_dim is None else self.coding_dim)
+
+
+@dataclass(frozen=True)
+class OverheadEstimate:
+    """Deterministic estimate of compute and memory costs for pooled coding.
+
+    Fields
+    - conv_macs: baseline Conv2d MACs per forward pass.
+    - coding_macs_total: total MACs added by pooled coding.
+    - coding_macs_breakdown: dict with keys ``pool``, ``projection``,
+      ``codebook`` (distance + assignment), and ``softmax`` (small terms).
+    - param_overhead: number of parameters introduced by the wrapper
+      (codebook centers + optional projection).
+    - activation_overhead_bytes: estimated transient activation bytes in the
+      coding path (pooled vector and soft-code; float elements only).
+    - overhead_ratio: ``coding_macs_total / conv_macs`` when ``conv_macs > 0``
+      else ``None``.
+    """
+
+    conv_macs: int
+    coding_macs_total: int
+    coding_macs_breakdown: dict[str, int]
+    param_overhead: int
+    activation_overhead_bytes: int
+    overhead_ratio: float | None
+
+
+def estimate_pooled_conv_overhead(cfg: PooledConvConfig) -> OverheadEstimate:
+    """Estimate compute and memory overhead for pooled Conv2d coding.
+
+    The estimator uses simple closed-form counts for MACs and activation sizes
+    that match the pooled-coding path in ``CodingConv2d``. See module docstring
+    for conventions and caveats.
+    """
+
+    B = int(cfg.batch)
+    Cin = int(cfg.in_channels)
+    Cout = int(cfg.out_channels)
+    kH, kW = (int(cfg.kernel_size[0]), int(cfg.kernel_size[1]))
+    H_out, W_out = cfg.output_size()
+    D = int(cfg.resolved_coding_dim())
+    K = int(cfg.K)
+    bytes_per = max(1, int(cfg.dtype_bytes))
+
+    # Baseline conv MACs (multiply-accumulate count)
+    conv_macs = B * H_out * W_out * Cout * (Cin * kH * kW)
+
+    # Coding computations
+    # 1) Global average pool over spatial dims: sum H*W then divide (per channel)
+    pool_macs = B * Cout * H_out * W_out
+
+    # 2) Optional linear projection when D < Cout
+    proj_macs = 0
+    proj_params = 0
+    if Cout > D:
+        # Matrix-vector multiply per sample: (Cout x D) @ (B x Cout)
+        proj_macs = B * Cout * D
+        proj_params = Cout * D
+
+    # 3) Codebook distances and soft assignment on (B, D) against K centers
+    #    Approximate as ~ 2 * B * K * D MACs (diff + square + sum)
+    codebook_macs = 2 * B * K * D
+    # Small softmax/logsumexp overhead per sample; keep it simple
+    softmax_macs = B * K
+
+    coding_macs_total = pool_macs + proj_macs + codebook_macs + softmax_macs
+
+    # Parameter overhead: K * D codebook centers + optional projection
+    param_overhead = K * D + proj_params
+
+    # Transient activation bytes (float elements): pooled (B, D) + soft (B, K)
+    activation_overhead_bytes = (B * D + B * K) * bytes_per
+
+    overhead_ratio = float(coding_macs_total) / float(conv_macs) if conv_macs > 0 else None
+
+    return OverheadEstimate(
+        conv_macs=int(conv_macs),
+        coding_macs_total=int(coding_macs_total),
+        coding_macs_breakdown={
+            "pool": int(pool_macs),
+            "projection": int(proj_macs),
+            "codebook": int(codebook_macs),
+            "softmax": int(softmax_macs),
+        },
+        param_overhead=int(param_overhead),
+        activation_overhead_bytes=int(activation_overhead_bytes),
+        overhead_ratio=None if overhead_ratio is None else float(overhead_ratio),
+    )
+
+
+__all__ = [
+    "OverheadEstimate",
+    "PooledConvConfig",
+    "estimate_pooled_conv_overhead",
+]