nervecode 0.1.0__py3-none-any.whl

nervecode/layers/linear.py

@@ -0,0 +1,176 @@
+"""Coding wrapper for ``torch.nn.Linear``.
+
+This module provides ``CodingLinear`` — the first production wrapper that
+observes a linear layer's output, reduces it (identity in the MVP), and
+computes a coding trace via a learnable codebook and a soft assignment engine.
+
+Observe-only contract: the wrapper never modifies the value returned to the
+caller; coding runs side-by-side and updates an internal trace cache.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .base import BaseCodingWrapper, ReducerFn
+
+# Import torch lazily but with typing support. Keep runtime import optional so the
+# package remains importable in environments without torch (tests will skip).
+try:  # pragma: no cover - exercised in torch-enabled environments
+    import torch
+    from torch import nn
+except Exception:  # pragma: no cover - allow import without torch
+    from typing import Any
+    from typing import cast as _cast
+
+    torch = _cast(Any, None)
+    nn = _cast(Any, None)
+
+# Local core imports guarded by torch availability at use sites.
+from nervecode.core import CodingTrace
+from nervecode.core.assignment import SoftAssignment
+from nervecode.core.codebook import Codebook
+
+__all__ = ["CodingLinear"]
+
+
+class CodingLinear(nn.Module, BaseCodingWrapper):
+    """Observe-only wrapper around ``nn.Linear``.
+
+    The wrapper delegates computation to the underlying ``nn.Linear`` and, when
+    coding is active, computes a ``CodingTrace`` over the visible activation.
+    The visible forward output is never altered.
+
+    Parameters
+    - layer: The ``nn.Linear`` module to wrap.
+    - K: Number of codebook centers. Defaults to 16.
+    - coding_dim: Target coding dimension ``D``. When ``None`` (default),
+      uses ``layer.out_features``. If provided, must satisfy
+      ``1 <= coding_dim <= layer.out_features``. A learned linear projection
+      reducer is created automatically when ``layer.out_features > coding_dim``;
+      otherwise the identity reduction is used.
+    - reducer: Optional reducer callable. Overrides ``coding_dim`` behavior
+      when provided. Defaults to identity (no reduction).
+    - temperature: Soft-assignment temperature. Higher is softer.
+    - eps: Numerical epsilon to clamp divisions and logs.
+
+    Notes
+    - MVP uses identity reduction; when a lower coding dimension is desired, a
+      learned projection reducer will be added in a follow-up task.
+    - ``repr(wrapper)`` includes ``code_dim``, codebook size ``K``, reducer type,
+      and whether coding is enabled to make printed models self-descriptive.
+    """
+
+    def __init__(
+        self,
+        layer: Any,  # use Any to avoid a hard torch dependency in type hints
+        *,
+        K: int = 16,
+        coding_dim: int | None = None,
+        reducer: ReducerFn | None = None,
+        temperature: float = 1.0,
+        eps: float = 1e-12,
+    ) -> None:
+        if torch is None:  # pragma: no cover - defensive
+            raise RuntimeError("CodingLinear requires PyTorch to be installed")
+
+        nn.Module.__init__(self)
+        BaseCodingWrapper.__init__(self, reducer=reducer)
+
+        if not isinstance(layer, nn.Linear):
+            raise TypeError("layer must be an instance of torch.nn.Linear")
+
+        self.layer: Any = layer
+
+        # Decide coding dimension and set up reducer if not provided.
+        out_features = int(self.layer.out_features)
+        if coding_dim is None:
+            code_dim = out_features
+        else:
+            code_dim = int(coding_dim)
+            if code_dim <= 0:
+                raise ValueError("coding_dim must be >= 1")
+            if code_dim > out_features:
+                raise ValueError("coding_dim must be <= layer.out_features for reduction")
+
+        # Optional learned projection reducer when output width exceeds coding dim.
+        self._projection: Any | None = None
+        if reducer is None and code_dim < out_features:
+            # Register a submodule to learn the reduction; keep observe-only semantics
+            # by applying it only for coding. Bias is unnecessary for dimension change.
+            self._projection = nn.Linear(out_features, code_dim, bias=False)
+
+            def _project_reducer(x: Any) -> tuple[Any, dict[str, Any]]:
+                y_proj = self._projection(x)  # type: ignore[misc]
+                meta = {"method": "linear_projection", "in": out_features, "out": code_dim}
+                return y_proj, meta
+
+            self.set_reducer(_project_reducer)
+        # else: keep provided reducer or identity_reduction
+
+        # Codebook operates over the last (feature) dimension using the decided code_dim.
+        self.codebook = Codebook(K=int(K), code_dim=code_dim)
+        self.assignment = SoftAssignment(temperature=float(temperature), eps=float(eps))
+
+    def forward(self, x: Any) -> Any:  # torch.Tensor in torch-enabled envs
+        """Return the wrapped layer's output and update the cached trace.
+
+        This method preserves the observe-only contract: the visible value is
+        the underlying layer output. For robust programmatic use that avoids
+        reading from mutable state, use ``forward_with_trace`` to retrieve the
+        explicit ``CodingTrace`` alongside the output.
+        """
+
+        y, _ = self.forward_with_trace(x)
+        return y
+
+    def forward_with_trace(self, x: Any) -> tuple[Any, CodingTrace | None]:
+        """Return the wrapped output and, when active, the explicit trace.
+
+        When coding is disabled or bypassed, returns ``(y, None)`` while
+        clearing any previously cached trace to avoid confusion. When active,
+        computes a fresh ``CodingTrace``, updates the internal cache, and
+        returns ``(y, trace)``.
+        """
+
+        y = self.layer(x)
+
+        # Fast path: skip all coding work when disabled or bypassed.
+        if not self.is_coding_active:
+            # Clear stale trace to avoid confusion when toggling modes.
+            self.clear_last_trace()
+            return y, None
+
+        # Apply the configured reducer (identity by default in MVP).
+        reduced, reduction_meta = self.maybe_reduce(y)
+
+        # Compute soft assignments and gather intermediates for the trace.
+        soft, inter = self.assignment(reduced, self.codebook)
+
+        trace = CodingTrace(
+            reduced=reduced,
+            reduction_meta=reduction_meta,
+            nearest_center_distances=inter["nearest_center_distances"],
+            chosen_center_indices=inter["chosen_center_indices"],
+            commitment_distances=inter["commitment_distances"],
+            soft_code=soft,
+        )
+        self.update_last_trace(trace)
+        return y, trace
+
+    def extra_repr(self) -> str:  # pragma: no cover - cosmetic
+        """Human-friendly summary shown inside ``nn.Module.__repr__``.
+
+        Includes the coding-space dimension (``code_dim``), codebook size
+        (``K``), the configured reducer type, and whether coding is currently
+        enabled. Also echoes the wrapped linear layer's ``out_features`` for
+        quick visual alignment.
+        """
+
+        status = "on" if self.coding_enabled else "off"
+        return (
+            f"layer=Linear(out_features={int(self.layer.out_features)}), "
+            f"code_dim={int(self.codebook.code_dim)}, "
+            f"K={int(self.codebook.K)}, reducer={self.get_reducer().__name__}, "
+            f"coding={status}"
+        )

nervecode/layers/reducers.py

@@ -0,0 +1,80 @@
+"""Experimental reducers for convolutional wrappers.
+
+This module provides optional, opt-in reducers that can be passed to
+``CodingConv2d`` (or any wrapper built on ``BaseCodingWrapper``) to prototype
+alternative spatial reductions beyond the default global average pooling.
+
+The helpers return callables that match ``ReducerFn`` so they can be supplied
+via the existing ``reducer=...`` parameter without changing stable defaults.
+
+Notes
+- These reducers are experimental and off by default. They preserve the
+  observe-only contract: reductions are applied only on the coding path.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .base import ReducerFn
+
+__all__ = [
+    "global_max_pool2d",
+    "spatial_tokens_identity_channels",
+]
+
+
+def global_max_pool2d() -> ReducerFn:
+    """Return a reducer that applies global max pooling over H and W.
+
+    Expects a 4D tensor with NCHW layout (``(B, C, H, W)``) and returns a
+    pooled tensor of shape ``(B, C)``. Metadata records the reduction method
+    and axes.
+    """
+
+    def _reduce(x: Any) -> tuple[Any, dict[str, Any]]:
+        # Require a 4D tensor-like object with an `amax` reduction
+        if getattr(x, "ndim", 0) != 4 or not hasattr(x, "amax"):
+            raise ValueError("global_max_pool2d expects a 4D NCHW tensor from Conv2d (B, C, H, W)")
+        # Max over spatial dims (H, W) -> (B, C)
+        reduced = x.amax(dim=(-1, -2))
+        meta = {
+            "method": "global_max_pool2d",
+            "spatial_reduction": True,
+            "reduction_axes": [-2, -1],  # H, W in NCHW
+        }
+        return reduced, meta
+
+    # Return the callable reducer
+    return _reduce
+
+
+def spatial_tokens_identity_channels() -> ReducerFn:
+    """Return a reducer that exposes per-spatial-position tokens with C last.
+
+    Transforms a standard Conv2d output ``(B, C, H, W)`` into a token-like
+    layout ``(B, H, W, C)`` without reducing channels. This allows computing
+    assignments per spatial position while keeping observe-only semantics.
+
+    The metadata clarifies that the view is token-like (no spatial reduction).
+    """
+
+    def _reduce(x: Any) -> tuple[Any, dict[str, Any]]:
+        # Require a 4D tensor-like object with `permute` and `contiguous`
+        if getattr(x, "ndim", 0) != 4 or not hasattr(x, "permute"):
+            raise ValueError(
+                "spatial_tokens_identity_channels expects a 4D NCHW tensor (B, C, H, W)"
+            )
+        # Permute to BHWC so that the feature dimension is last for coding.
+        reduced = x.permute(0, 2, 3, 1).contiguous()
+        meta = {
+            "method": "spatial_tokens_identity",
+            "token_like": True,
+            "channels_last": True,
+            "spatial_reduction": False,
+            "source_layout": "NCHW",
+            "target_layout": "NHWC",
+        }
+        return reduced, meta
+
+    return _reduce

nervecode/layers/wrap.py

@@ -0,0 +1,223 @@
+"""Model instrumentation helpers for coding wrappers.
+
+This module provides a small `wrap()` function that can instrument a PyTorch
+`nn.Module` in-place based on either:
+
+- explicit dotted module names (e.g., ["encoder.fc", "head.1"]), or
+- supported selection shortcuts (MVP: "all_linear").
+
+Only supported layer types are wrapped; others are left unchanged to preserve
+fail-open behavior. The function returns the same model instance for
+convenience.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+try:  # Keep runtime import optional for environments without torch
+    import torch
+    from torch import nn
+except Exception:  # pragma: no cover - import-time optionality
+    from typing import Any
+    from typing import cast as _cast
+
+    torch = _cast(Any, None)
+    nn = _cast(Any, None)
+
+from .conv import CodingConv2d
+from .linear import CodingLinear
+
+__all__ = ["wrap"]
+
+
+def wrap(
+    model: Any,
+    *,
+    layers: str | Iterable[str] = "all_linear",
+    K: int | None = None,
+    coding_dim: int | None = None,
+    temperature: float = 1.0,
+    eps: float = 1e-12,
+) -> Any:
+    """Instrument ``model`` in-place with coding wrappers.
+
+    Parameters
+    - model: A PyTorch ``nn.Module`` to instrument. The function modifies the
+      module tree in-place and returns the same instance for convenience.
+    - layers: Either a selection shortcut string (MVP supports only
+      ``"all_linear"``) or an iterable of explicit dotted module names to
+      wrap. Names must match those returned by ``model.named_modules()``.
+
+    Behavior
+    - Only supported layer types are wrapped. Currently this includes
+      ``nn.Linear`` via ``CodingLinear``. Unsupported selections are ignored to
+      preserve fail-open behavior.
+    - Missing names are ignored.
+    - On environments without PyTorch, the function is a no-op and returns the
+      model unchanged.
+    """
+
+    if nn is None:  # pragma: no cover - defensive; torch is a dependency in tests
+        return model
+
+    # Validate basic interface; fall back to no-op when model is not an nn.Module
+    if not isinstance(model, nn.Module):
+        return model
+
+    # Resolve target names to instrument
+    if isinstance(layers, str):
+        target_names = _select_by_shortcut(model, layers)
+    else:
+        # Deduplicate while preserving order
+        seen: set[str] = set()
+        target_names = []
+        for name in layers:
+            if not isinstance(name, str):
+                continue
+            if name not in seen:
+                seen.add(name)
+                target_names.append(name)
+
+    # Map types to wrapper factories; extend as new wrappers are added.
+    def _wrap_module(m: nn.Module) -> nn.Module | None:
+        if isinstance(m, nn.Linear):
+            kwargs: dict[str, Any] = {
+                "temperature": float(temperature),
+                "eps": float(eps),
+            }
+            if K is not None:
+                kwargs["K"] = int(K)
+            if coding_dim is not None:
+                # Clamp coding_dim per layer to avoid invalid configuration
+                try:
+                    cd_req = int(coding_dim)
+                    cd_eff = max(1, min(int(m.out_features), cd_req))
+                except Exception:
+                    cd_eff = int(m.out_features)
+                kwargs["coding_dim"] = cd_eff
+            # Defensive: if downstream validation still rejects coding_dim,
+            # fall back to using the layer's out_features to avoid crashes.
+            try:
+                return CodingLinear(m, **kwargs)
+            except Exception as exc:
+                msg = str(exc)
+                if "coding_dim must be <= layer.out_features" in msg:
+                    safe_kwargs = dict(kwargs)
+                    safe_kwargs["coding_dim"] = int(getattr(m, "out_features", 1) or 1)
+                    return CodingLinear(m, **safe_kwargs)
+                raise
+        if isinstance(m, nn.Conv2d):
+            kwargs_c: dict[str, Any] = {
+                "temperature": float(temperature),
+                "eps": float(eps),
+            }
+            if K is not None:
+                kwargs_c["K"] = int(K)
+            if coding_dim is not None:
+                try:
+                    cd_req = int(coding_dim)
+                    cd_eff = max(1, min(int(m.out_channels), cd_req))
+                except Exception:
+                    cd_eff = int(m.out_channels)
+                kwargs_c["coding_dim"] = cd_eff
+            try:
+                return CodingConv2d(m, **kwargs_c)
+            except Exception as exc:
+                msg = str(exc)
+                if "coding_dim must be <=" in msg:
+                    safe_kwargs_c = dict(kwargs_c)
+                    safe_kwargs_c["coding_dim"] = int(getattr(m, "out_channels", 1) or 1)
+                    return CodingConv2d(m, **safe_kwargs_c)
+                raise
+        return None
+
+    # Instrument selected modules by replacing them on their parent.
+    for name in target_names:
+        try:
+            # Skip the top-level module itself
+            if name == "":
+                continue
+            # get_submodule raises on unknown names; ignore gracefully
+            sub = model.get_submodule(name)
+        except Exception:
+            continue
+        wrapped = _wrap_module(sub)
+        if wrapped is None:
+            continue  # unsupported type
+        parent_path, leaf = _split_parent(name)
+        try:
+            parent = model.get_submodule(parent_path) if parent_path else model
+        except Exception:
+            continue
+        # Replace on the parent; handle non-identifier names (e.g., Sequential indices)
+        if leaf.isidentifier():
+            setattr(parent, leaf, wrapped)
+        else:
+            # Fall back to the internal registry for names like "1" in Sequential
+            parent._modules[leaf] = wrapped
+
+    return model
+
+
+def _split_parent(path: str) -> tuple[str, str]:
+    """Return (parent_path, leaf) for a dotted module path.
+
+    Examples
+    - "fc1" -> ("", "fc1")
+    - "block.1" -> ("block", "1")
+    - "encoder.layer.0.attn" -> ("encoder.layer.0", "attn")
+    """
+
+    if "." not in path:
+        return "", path
+    parent, leaf = path.rsplit(".", 1)
+    return parent, leaf
+
+
+def _select_by_shortcut(model: nn.Module, key: str) -> list[str]:
+    """Return module names selected by a supported shortcut key.
+
+    Supported keys (MVP):
+    - "all_linear": select every ``nn.Linear`` in ``model``.
+    - "all_conv": select every ``nn.Conv2d`` in ``model``.
+    Unrecognized keys select nothing (fail-open).
+    """
+
+    key_norm = key.strip().lower()
+
+    # all_* shortcuts: filter across named_modules
+    if key_norm == "all_linear":
+        return [name for name, m in model.named_modules() if isinstance(m, nn.Linear)]
+    if key_norm == "all_conv":
+        return [name for name, m in model.named_modules() if isinstance(m, nn.Conv2d)]
+
+    # first/last convenience: select a single Linear
+    if key_norm == "first_linear":
+        for name, m in model.named_modules():
+            if isinstance(m, nn.Linear) and name:
+                return [name]
+        return []
+    if key_norm == "last_linear":
+        last: str | None = None
+        for name, m in model.named_modules():
+            if isinstance(m, nn.Linear) and name:
+                last = name
+        return [last] if last else []
+
+    # first/last convenience for Conv2d
+    if key_norm == "first_conv":
+        for name, m in model.named_modules():
+            if isinstance(m, nn.Conv2d) and name:
+                return [name]
+        return []
+    if key_norm == "last_conv":
+        last_c: str | None = None
+        for name, m in model.named_modules():
+            if isinstance(m, nn.Conv2d) and name:
+                last_c = name
+        return [last_c] if last_c else []
+
+    # Unknown shortcut -> empty selection (fail-open)
+    return []

nervecode/scoring/__init__.py

@@ -0,0 +1,20 @@
+"""Scoring, aggregation, and calibration.
+
+Scoring functions turn per-layer codes and traces into useful signals such as
+surprise estimates. Calibration utilities live here as well.
+"""
+
+from __future__ import annotations
+
+from .aggregator import max_surprise, mean_surprise, weighted_surprise
+from .calibrator import CalibrationState, EmpiricalPercentileCalibrator
+from .types import AggregatedSurprise
+
+__all__: list[str] = [
+    "AggregatedSurprise",
+    "CalibrationState",
+    "EmpiricalPercentileCalibrator",
+    "max_surprise",
+    "mean_surprise",
+    "weighted_surprise",
+]