nervecode 0.1.0__py3-none-any.whl

nervecode/__init__.py

@@ -0,0 +1,415 @@
+"""Nervecode package root and public API stubs.
+
+The real implementation will land incrementally following the MVP plan. For now,
+we expose the intended top-level API so users and tooling can rely on the package
+surface existing from the start.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator
+from contextlib import ExitStack, contextmanager, suppress
+from typing import Any
+
+# Import base wrapper type for lightweight wrapper tracking without a hard
+# torch dependency at import time.
+try:  # pragma: no cover - available in normal package use
+    from nervecode.layers.base import BaseCodingWrapper
+except Exception:  # pragma: no cover - fallback for environments without package context
+    BaseCodingWrapper = object  # type: ignore[misc,assignment]
+
+
+class WrappedModel:
+    """Thin container around a base model that tracks coding layers.
+
+    Goals
+    - Preserve the original model API: attribute access and ``__call__``
+      delegate to the wrapped model so existing code keeps working.
+    - Track inserted coding layers (wrappers) to support convenience
+      iterators and aggregations implemented in subsequent tasks.
+    """
+
+    # Keep the constructor torch-agnostic and accept any object. When the
+    # object is a PyTorch ``nn.Module``, wrapper discovery uses
+    # ``named_modules()``; otherwise, tracking stays empty.
+    def __init__(self, model: object) -> None:
+        # Store the wrapped model reference for external access and tests.
+        self._model = model
+        # Internal registry of discovered coding wrappers keyed by dotted name.
+        self._wrapped_layers: dict[str, BaseCodingWrapper] = {}
+        # Convenience caches populated on forward() when coding is active.
+        # Latest visible model output from the most recent forward.
+        self._last_output: Any | None = None
+        # Latest per-layer traces keyed by dotted module name.
+        # Values are intentionally typed as Any to keep this module torch-agnostic;
+        # concrete type is nervecode.core.CodingTrace in torch-enabled envs.
+        self._last_layer_traces: dict[str, Any] = {}
+        # Optional calibrator attached after a calibration run.
+        self._calibrator: Any | None = None
+        self._refresh_wrapped_layers()
+
+    # --- Delegation to preserve the base model API -----------------------------
+    def __getattr__(self, name: str) -> Any:  # pragma: no cover - exercised via tests
+        # Delegate attribute access to the wrapped model when not found on self.
+        try:
+            return getattr(self._model, name)
+        except Exception as exc:  # propagate standard AttributeError semantics
+            raise exc
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Delegate calls to the wrapped model's ``__call__``.
+
+        This keeps the container usable anywhere the base model is expected.
+        The explicit ``forward()`` behavior is implemented in a later task.
+        """
+
+        model = self._model
+        # Fast path: call the model if it is callable.
+        if callable(model):
+            return model(*args, **kwargs)
+        # Fallback: attempt a direct ``forward`` attribute.
+        fwd = getattr(model, "forward", None)
+        if callable(fwd):
+            return fwd(*args, **kwargs)
+        raise TypeError("Wrapped base model is not callable and has no forward()")
+
+    # --- Common nn.Module-style helpers that preserve chaining semantics -------
+    def train(self, mode: bool = True) -> WrappedModel:  # pragma: no cover - exercised indirectly
+        """Set underlying model to train/eval mode and return self for chaining.
+
+        Mirrors ``torch.nn.Module.train`` semantics so that patterns like
+        ``WrappedModel(model).train()`` keep returning a wrapper rather than the
+        bare model (which would drop wrapper methods).
+        """
+
+        m = getattr(self._model, "train", None)
+        if callable(m):
+            with suppress(Exception):
+                m(bool(mode))
+        return self
+
+    def eval(self) -> WrappedModel:  # pragma: no cover - exercised indirectly
+        """Switch to evaluation mode and return self for chaining."""
+        return self.train(False)
+
+    def forward(self, *args: Any, **kwargs: Any) -> object:
+        """Run a forward pass through the wrapped base model.
+
+        Returns the base model output. In the MVP, this will also collect coding
+        statistics when enabled.
+        """
+        # Ensure our registry reflects any instrumentation that may have been
+        # applied after construction.
+        self._refresh_wrapped_layers()
+
+        # Clear last-forward convenience caches before computing a new result.
+        self._last_output = None
+        self._last_layer_traces = {}
+
+        # Delegate to the underlying model using its normal call semantics.
+        model = self._model
+        if callable(model):
+            out = model(*args, **kwargs)
+        else:
+            fwd = getattr(model, "forward", None)
+            if callable(fwd):
+                out = fwd(*args, **kwargs)
+            else:
+                raise TypeError("Wrapped base model is not callable and has no forward()")
+
+        # Update convenience caches: store output and collect latest per-layer traces
+        # from wrapped modules when coding is enabled at the layer level.
+        self._last_output = out
+        try:
+            traces: dict[str, Any] = {}
+            for name, wrapper in self._wrapped_layers.items():
+                # Each wrapper manages its own fail-open behavior and cache.
+                trace = getattr(wrapper, "last_trace", None)
+                if trace is not None:
+                    traces[name] = trace
+            self._last_layer_traces = traces
+        except Exception:
+            # Fail-open: keep caches best-effort without affecting user forwards.
+            self._last_layer_traces = {}
+
+        return out
+
+    # --- Model-level fail-open controls ----------------------------------------
+    def enable_coding(self) -> None:
+        """Enable coding on all discovered wrapped layers.
+
+        Safe to call at any time; silently does nothing when no wrappers are
+        present or when the wrapped object is not a PyTorch module.
+        """
+
+        self._refresh_wrapped_layers()
+        try:
+            for wrapper in self._wrapped_layers.values():
+                wrapper.enable_coding()
+        except Exception:
+            # Fail-open: best-effort delegation only
+            return
+
+    def disable_coding(self) -> None:
+        """Disable coding on all discovered wrapped layers (fail-open).
+
+        Clearing any cached per-layer traces ensures that callers observing
+        model-level caches see the disabled state immediately, without having
+        to run another forward first.
+        """
+
+        self._refresh_wrapped_layers()
+        try:
+            for wrapper in self._wrapped_layers.values():
+                wrapper.disable_coding()
+                # Ensure stale traces are not exposed after disabling.
+                if hasattr(wrapper, "clear_last_trace"):
+                    wrapper.clear_last_trace()
+        except Exception:
+            # Fail-open: best-effort delegation only
+            pass
+        # Clear model-level convenience cache as well.
+        self._last_layer_traces = {}
+
+    @contextmanager
+    def bypass(self) -> Iterator[None]:
+        """Temporarily bypass coding across all wrapped layers.
+
+        The context nests correctly and delegates to every discovered wrapper's
+        own bypass context. Safe when no wrappers are present.
+        """
+
+        self._refresh_wrapped_layers()
+        with ExitStack() as stack:
+            for wrapper in self._wrapped_layers.values():
+                # Each wrapper manages its own nesting depth.
+                stack.enter_context(wrapper.bypass())
+            yield
+
+    # Training-time auxiliary loss used to learn codebooks.
+    def coding_loss(self) -> Any:
+        """Compute coding loss from the latest per-layer traces.
+
+        Uses the most recent per-layer traces collected during ``forward`` and
+        computes a scalar training-time loss via ``nervecode.training.CodingLoss``.
+        When no traces are available, raises a helpful error explaining how to
+        proceed (e.g., instrument the model, enable coding, and run a forward
+        pass, or use the explicit trace path and call ``CodingLoss`` directly).
+        """
+
+        # Lazy import to keep package importable without torch in non-training contexts
+        try:
+            from nervecode.training import CodingLoss
+        except Exception as exc:  # pragma: no cover - defensive
+            raise RuntimeError(
+                "CodingLoss is unavailable. Ensure PyTorch and training components are installed."
+            ) from exc
+
+        traces = getattr(self, "_last_layer_traces", {}) or {}
+        if not traces:
+            # Diagnose common causes and provide actionable guidance.
+            wrappers = getattr(self, "_wrapped_layers", {}) or {}
+            if not wrappers:
+                raise RuntimeError(
+                    "No coding layers discovered on this model. Instrument your model "
+                    "(e.g., nervecode.layers.wrap.wrap(model, layers='all_linear')) "
+                    "before calling WrappedModel.coding_loss()."
+                )
+            any_enabled = any(getattr(w, "coding_enabled", False) for w in wrappers.values())
+            if not any_enabled:
+                raise RuntimeError(
+                    "Coding is disabled on all wrapped layers. Call wrapped.enable_coding() "
+                    "and run a forward pass before calling WrappedModel.coding_loss()."
+                )
+            raise RuntimeError(
+                "No per-layer traces are available. Run wrapped.forward(...) first to "
+                "populate the latest traces, or use forward_with_trace(...) and pass explicit "
+                "traces to nervecode.training.CodingLoss."
+            )
+
+        loss_mod = CodingLoss()
+        return loss_mod(traces)
+
+    # Calibrate surprise thresholds/percentiles on in-distribution data.
+    def calibrate(self, data_loader: Iterable[Any]) -> None:
+        """Calibrate empirical percentiles on aggregated surprise.
+
+        Runs the wrapped model over a data loader (held-out in-distribution),
+        collects per-sample aggregated surprise values, and fits an
+        EmpiricalPercentileCalibrator. The fitted calibrator is attached to the
+        instance for later use (e.g., percentiles, is_ood) via ``self._calibrator``.
+
+        Notes
+        - The data loader may yield either ``inputs`` or ``(inputs, targets)``.
+        - Coding must be enabled on wrapped layers for traces to be produced.
+        """
+
+        try:
+            from nervecode.scoring import EmpiricalPercentileCalibrator, mean_surprise
+        except Exception as exc:  # pragma: no cover - defensive
+            raise RuntimeError(
+                "Calibration utilities are unavailable. Ensure nervecode.scoring is importable."
+            ) from exc
+
+        # Collect per-sample surprise vectors across the loader
+        scores: list[Any] = []
+        # Prefer torch.no_grad if available, but avoid hard dependency at import time.
+        try:
+            import torch
+
+            no_grad_ctx: Any = torch.no_grad()
+        except Exception:  # pragma: no cover - environments without torch
+
+            class _NullCtx:
+                def __enter__(self, *a: Any, **k: Any) -> None:
+                    return None
+
+                def __exit__(self, *a: Any, **k: Any) -> None:
+                    return None
+
+            no_grad_ctx = _NullCtx()
+
+        with no_grad_ctx:
+            for batch in data_loader:
+                # Common patterns: (inputs, targets) or inputs directly
+                if isinstance(batch, (tuple, list)) and len(batch) >= 1:
+                    inputs = batch[0]
+                elif isinstance(batch, dict) and "inputs" in batch:
+                    inputs = batch["inputs"]
+                else:
+                    inputs = batch
+
+                _ = self.forward(inputs)
+                agg = self.surprise()
+                if agg is None:
+                    # Fall back to explicit aggregation over the latest traces
+                    traces = getattr(self, "_last_layer_traces", {}) or {}
+                    agg = mean_surprise(traces)
+                if agg is not None and hasattr(agg, "surprise"):
+                    scores.append(agg.surprise)
+
+        if not scores:
+            raise RuntimeError(
+                "Calibration collected no surprise scores; ensure coding is enabled."
+            )
+
+        # Fit calibrator on concatenated scores when possible (torch),
+        # otherwise fit per-chunk (rare path when tensor ops are unavailable).
+        try:
+            import torch
+
+            vec = torch.cat([s.detach().to(device="cpu").reshape(-1) for s in scores], dim=0)
+            calib = EmpiricalPercentileCalibrator(threshold_quantiles=(0.95,))
+            _ = calib.fit(vec, aggregation="mean")
+        except Exception:  # pragma: no cover - defensive minimal fallback
+            calib = EmpiricalPercentileCalibrator(threshold_quantiles=(0.95,))
+            # Let calibrator normalize iterable inputs internally
+            flat: list[float] = []
+            for s in scores:
+                with suppress(Exception):
+                    resh = getattr(s, "reshape", None)
+                    if callable(resh):
+                        arr = resh(-1).tolist()
+                    else:
+                        arr = s.tolist() if hasattr(s, "tolist") else [s]
+                    flat.extend([float(v) for v in list(arr)])
+            _ = calib.fit(flat, aggregation="mean")
+
+        self._calibrator = calib
+        return None
+
+    # Inference-time surprise signal (aggregated over wrapped layers).
+    def surprise(self) -> object:
+        """Return the latest aggregated surprise across wrapped layers.
+
+        This convenience method aggregates the most recent per-layer traces
+        collected during a standard ``forward`` into a per-sample surprise
+        signal using the default mean strategy. When no valid per-layer
+        signals are available, returns ``None``.
+        """
+
+        try:
+            from nervecode.scoring import mean_surprise  # lazy import
+        except Exception:
+            return None
+
+        traces = getattr(self, "_last_layer_traces", {}) or {}
+        return mean_surprise(traces)
+
+    # Explicit trace-returning path for robust integrations.
+    def forward_with_trace(self, *args: Any, **kwargs: Any) -> tuple[object, dict[str, Any]]:
+        """Return (model_output, per-layer_traces) for a single forward.
+
+        The first element mirrors the wrapped model's normal return value. The
+        second is a mapping from dotted layer names to ``CodingTrace`` objects
+        produced during this forward by wrapped layers (those with coding
+        active). When coding is disabled, the mapping is empty.
+        """
+
+        out = self.forward(*args, **kwargs)
+        traces = dict(getattr(self, "_last_layer_traces", {}) or {})
+        return out, traces
+
+    # --- Internal: discover coding wrappers on the model -----------------------
+    def _refresh_wrapped_layers(self) -> None:
+        """Discover and cache coding wrappers present on the wrapped model.
+
+        Populates ``_wrapped_layers`` with a mapping from dotted module names to
+        wrapper instances. Non-PyTorch objects result in an empty registry.
+        The method never raises; it is safe to call after in-place
+        instrumentation.
+        """
+
+        registry: dict[str, BaseCodingWrapper] = {}
+        model = self._model
+        try:
+            named_modules = getattr(model, "named_modules", None)
+            if callable(named_modules):
+                for name, module in named_modules():
+                    # isinstance check remains torch-agnostic because
+                    # BaseCodingWrapper is torch-free.
+                    if isinstance(module, BaseCodingWrapper):
+                        registry[str(name)] = module
+        except Exception:
+            # Fail-open: keep tracking best-effort only.
+            registry = {}
+        self._wrapped_layers = registry
+
+
+def wrap(model: object, *, layers: str | Iterable[str] = "all_linear") -> WrappedModel:
+    """Wrap a PyTorch module to produce coding-based surprise signals.
+
+    Parameters
+    - model: Base `torch.nn.Module` to wrap.
+    - layers: Strategy or explicit layer names to wrap. The MVP will start with
+      the convenience value "all_linear".
+    """
+    # Import the in-place instrumentation helper lazily to keep import-time
+    # dependencies minimal in environments without torch.
+    try:
+        from nervecode.layers.wrap import wrap as _wrap_layers
+    except Exception as exc:  # pragma: no cover - defensive
+        raise RuntimeError(
+            "Layer instrumentation is unavailable. Import nervecode.layers.wrap directly."
+        ) from exc
+
+    try:
+        _wrap_layers(model, layers=layers)
+    except Exception:
+        # Fail-open: return a container even if instrumentation could not be applied.
+        # This preserves the attribute and call delegation behavior for non-torch objects.
+        return WrappedModel(model)
+
+    return WrappedModel(model)
+
+
+__all__ = [
+    "WrappedModel",
+    "wrap",
+]
+
+# Runtime-accessible package version
+try:  # pragma: no cover - trivial import
+    from ._version import __version__ as __version__
+except Exception:  # pragma: no cover - fallback if version file unavailable
+    __version__ = "0.0.0"

nervecode/_version.py

@@ -0,0 +1,10 @@
+"""Package version for Nervecode.
+
+The version is the single source of truth for both runtime access via
+``nervecode.__version__`` and build-time metadata via Hatch's version plugin.
+"""
+
+__all__ = ["__version__"]
+
+# Follow PEP 440; bump as releases evolve.
+__version__ = "0.1.0"

nervecode/core/__init__.py

@@ -0,0 +1,19 @@
+"""Core data model, primitives, and engines.
+
+This subpackage will host the core types, codebook modules, assignment
+algorithms, and trace representations used across Nervecode.
+"""
+
+from __future__ import annotations
+
+from .temperature import CosineTemperature, FixedTemperature, TemperatureSchedule
+from .trace import CodingTrace
+from .types import SoftCode
+
+__all__: list[str] = [
+    "CodingTrace",
+    "CosineTemperature",
+    "FixedTemperature",
+    "SoftCode",
+    "TemperatureSchedule",
+]

nervecode/core/assignment.py

@@ -0,0 +1,165 @@
+"""Soft assignment engine.
+
+This module computes squared Euclidean distances from reduced activations to a
+codebook, converts them into soft assignments over codes, and returns both a
+``SoftCode`` object and the intermediate tensors needed to construct a
+``CodingTrace`` without recomputing distances.
+
+Shapes follow the core convention:
+- Inputs have shape ``(..., D)`` where ``...`` are arbitrary leading dims.
+- Probabilities have shape ``(..., K)`` where ``K`` is the number of codes.
+- Per-position scalar values have shape matching the leading dims ``...``.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from .codebook import Codebook
+from .shapes import flatten_leading, unflatten_leading
+from .types import SoftCode
+
+__all__ = ["SoftAssignment"]
+
+
+class SoftAssignment(nn.Module):
+    """Compute soft assignments and return trace-ready intermediates.
+
+    Parameters
+    - temperature: Positive scalar controlling assignment sharpness. Larger
+      values produce softer assignments. This parameter is not learnable; use a
+      schedule or an external tensor if you need dynamic temperature.
+    - eps: Small positive constant used to clamp divisions and logarithms to
+      keep values finite.
+    """
+
+    def __init__(
+        self,
+        *,
+        temperature: float = 1.0,
+        eps: float = 1e-12,
+        beta_length: float = 1.0,
+        beta_entropy: float = 1.0,
+        beta_distance: float = 0.0,
+    ) -> None:
+        super().__init__()
+        t = float(temperature)
+        if not torch.isfinite(torch.tensor(t)) or t <= 0.0:
+            raise ValueError("temperature must be a positive finite float")
+        e = float(eps)
+        if not torch.isfinite(torch.tensor(e)) or e <= 0.0:
+            raise ValueError("eps must be a positive finite float")
+        # Explicit attribute annotation helps static type checkers understand
+        # the buffer's type (avoids Tensor|Module unions on .to/.clamp paths).
+        self._temperature: torch.Tensor
+        self.register_buffer("_temperature", torch.tensor(t), persistent=False)
+        self.eps: float = e
+        # Weights for the combined surprise signal S = bL * L + bH * H + bD * D_norm
+        self.beta_length: float = float(beta_length)
+        self.beta_entropy: float = float(beta_entropy)
+        self.beta_distance: float = float(beta_distance)
+
+    @property
+    def temperature(self) -> float:
+        return float(self._temperature.item())
+
+    def forward(
+        self,
+        reduced: torch.Tensor,
+        codebook: Codebook | torch.Tensor,
+    ) -> tuple[SoftCode, dict[str, torch.Tensor]]:
+        """Return soft code and intermediates for ``CodingTrace``.
+
+        Inputs
+        - reduced: Tensor of shape ``(..., D)`` representing reduced activations.
+        - codebook: ``Codebook`` module or a centers tensor of shape
+          ``(K, D)``.
+
+        Returns
+        - soft_code: ``SoftCode`` with probabilities ``(..., K)`` and basic
+          scalar statistics (best length, entropy, best indices, combined surprise).
+        - intermediates: dict with keys
+            - ``nearest_center_distances``: ``(...,)``
+            - ``chosen_center_indices``: integer ``(...,)``
+            - ``commitment_distances``: ``(...,)``
+          which are sufficient to construct a ``CodingTrace`` without
+          recomputation.
+        """
+
+        centers = codebook() if isinstance(codebook, Codebook) else codebook
+
+        if centers.ndim != 2:
+            raise ValueError("centers must have shape (K, D)")
+
+        x = reduced
+        if x.ndim < 1:
+            raise ValueError("reduced must have at least one dimension (..., D)")
+
+        x_flat, lead = flatten_leading(x)  # (N, D)
+        d = int(centers.shape[1])
+        if int(x_flat.shape[1]) != d:
+            raise ValueError(f"reduced last dim {int(x_flat.shape[1])} must match centers D={d}")
+
+        # Compute squared Euclidean distances using the expansion
+        # ||x - c||^2 = ||x||^2 + ||c||^2 - 2 x·c
+        x_norm = (x_flat * x_flat).sum(dim=1, keepdim=True)  # (N, 1)
+        c_norm = (centers * centers).sum(dim=1).unsqueeze(0)  # (1, K)
+        # x @ centers^T yields (N, K)
+        dot = x_flat @ centers.t()  # (N, K)
+        distances = x_norm + c_norm - 2.0 * dot  # (N, K)
+        # For numerical safety, clamp at zero from below (squared distances).
+        distances = distances.clamp_min(0.0)
+
+        # Convert distances into softmax probabilities over codes using
+        # negative distances as logits, scaled by temperature.
+        t = self._temperature.to(dtype=distances.dtype, device=distances.device)
+        logits = -distances / torch.clamp(t, min=self.eps)
+        probs_flat = torch.softmax(logits, dim=-1)
+
+        # Best indices from min distance (equivalently max probability).
+        best_idx_flat = torch.argmin(distances, dim=-1)  # (N,)
+        # Nearest-center distances and commitment distances per position.
+        nearest_flat = distances.gather(-1, best_idx_flat.unsqueeze(-1)).squeeze(-1)  # (N,)
+        commitment_flat = nearest_flat  # identical for squared Euclidean
+
+        # Information-theoretic scalars
+        eps_t = torch.tensor(self.eps, dtype=probs_flat.dtype, device=probs_flat.device)
+        best_prob_flat = probs_flat.gather(-1, best_idx_flat.unsqueeze(-1)).squeeze(-1)
+        best_length_flat = -(best_prob_flat + eps_t).log()
+        entropy_flat = -(probs_flat * (probs_flat + eps_t).log()).sum(dim=-1)
+        # Distance-based component: use a stable, monotonically increasing
+        # transform of the nearest-center squared Euclidean distance. ``log1p``
+        # reduces dynamic range without erasing separation. Avoid per-forward
+        # normalization to preserve absolute contrast between ID and OOD.
+        d_comp_flat = torch.log1p(nearest_flat)
+        # Weighted combination including optional distance contribution
+        combined_surprise_flat = (
+            self.beta_length * best_length_flat
+            + self.beta_entropy * entropy_flat
+            + self.beta_distance * d_comp_flat
+        )
+
+        # Restore original leading layout
+        probs = unflatten_leading(probs_flat, lead)  # (..., K)
+        best_idx = unflatten_leading(best_idx_flat, lead)  # (...,)
+        nearest = unflatten_leading(nearest_flat, lead)  # (...,)
+        commitment = unflatten_leading(commitment_flat, lead)  # (...,)
+        best_length = unflatten_leading(best_length_flat, lead)  # (...,)
+        entropy = unflatten_leading(entropy_flat, lead)  # (...,)
+        combined_surprise = unflatten_leading(combined_surprise_flat, lead)  # (...,)
+
+        soft = SoftCode(
+            probs=probs,
+            best_length=best_length,
+            entropy=entropy,
+            best_indices=best_idx,
+            combined_surprise=combined_surprise,
+        )
+
+        intermediates: dict[str, torch.Tensor] = {
+            "nearest_center_distances": nearest,
+            "chosen_center_indices": best_idx,
+            "commitment_distances": commitment,
+        }
+        return soft, intermediates