nervecode 0.1.0__py3-none-any.whl

nervecode/core/types.py

@@ -0,0 +1,116 @@
+"""Core type definitions.
+
+This module defines lightweight dataclasses used across the core coding
+pipeline. Shapes follow the convention that probability-like tensors carry
+arbitrary leading dimensions with a final code dimension of size ``K``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+__all__ = ["SoftCode"]
+
+
+@dataclass(frozen=True)
+class SoftCode:
+    """Soft assignment over a codebook.
+
+    The ``probs`` tensor encodes probabilities over ``K`` codes for each
+    position in the (possibly batched) input. It supports arbitrary leading
+    dimensions followed by a final code dimension of shape ``(..., K)``.
+
+    Fields
+    - `probs` (Tensor): probabilities over codes with shape ``(..., K)``.
+    - `best_length` (Tensor, optional): best-code codelength per position
+      with shape matching the leading dimensions ``...``.
+    - `entropy` (Tensor, optional): assignment entropy per position with
+      shape matching the leading dimensions ``...``.
+    - `best_indices` (Tensor, optional): argmax code indices per position with
+      shape matching the leading dimensions ``...``.
+    - `combined_surprise` (Tensor, optional): a scalar-like per-position
+      surprise that may combine multiple terms; shape matches ``...``.
+
+    All optional scalar-like fields must match the leading dimensions of
+    ``probs`` exactly (i.e., ``tensor.shape == probs.shape[:-1]``). When the
+    leading shape is empty, scalar-like fields must be 0-dim tensors.
+    """
+
+    probs: Any
+    best_length: Any | None = None
+    entropy: Any | None = None
+    best_indices: Any | None = None
+    combined_surprise: Any | None = None
+
+    def __post_init__(self) -> None:
+        from typing import Any
+        from typing import cast as _cast
+
+        try:
+            import torch  # runtime import; typing deferred
+        except Exception:  # pragma: no cover - torch absent in some test envs
+            torch = _cast(Any, None)
+
+        # Validate probs tensor
+        if torch is not None and not isinstance(
+            self.probs, torch.Tensor
+        ):  # pragma: no cover - defensive
+            raise TypeError("probs must be a torch.Tensor")
+        if getattr(self.probs, "ndim", None) is None or self.probs.ndim < 1:
+            raise ValueError("probs must have at least one dimension (..., K)")
+        k = int(self.probs.shape[-1])
+        if k <= 0:
+            raise ValueError("final code dimension K must be >= 1")
+
+        # Optional fields shape validation (only when torch is available)
+        if torch is None:  # pragma: no cover - validation relies on torch
+            return
+
+        lead_shape = tuple(int(s) for s in self.probs.shape[:-1])
+
+        def _check_shape(name: str, value: Any) -> None:
+            if not isinstance(value, torch.Tensor):
+                raise TypeError(f"{name} must be a torch.Tensor if provided")
+            if tuple(int(s) for s in value.shape) != lead_shape:
+                raise ValueError(
+                    f"{name} must have shape {lead_shape} to match probs.leading_shape"
+                )
+
+        if self.best_length is not None:
+            _check_shape("best_length", self.best_length)
+            # float-like recommended; do not hard-enforce dtype beyond non-bool
+            if self.best_length.dtype == torch.bool:  # pragma: no cover - defensive
+                raise TypeError("best_length must be numeric, not bool")
+
+        if self.entropy is not None:
+            _check_shape("entropy", self.entropy)
+            if self.entropy.dtype == torch.bool:  # pragma: no cover - defensive
+                raise TypeError("entropy must be numeric, not bool")
+
+        if self.combined_surprise is not None:
+            _check_shape("combined_surprise", self.combined_surprise)
+            if self.combined_surprise.dtype == torch.bool:  # pragma: no cover - defensive
+                raise TypeError("combined_surprise must be numeric, not bool")
+
+        if self.best_indices is not None:
+            _check_shape("best_indices", self.best_indices)
+            # Require integer dtype for indices
+            if torch.is_floating_point(self.best_indices) or self.best_indices.dtype == torch.bool:
+                raise TypeError("best_indices must use an integer dtype")
+
+    @property
+    def leading_shape(self) -> tuple[int, ...]:
+        """Return leading shape before the final code dimension.
+
+        For ``probs`` with shape ``(..., K)``, this returns ``...`` as a tuple.
+        """
+
+        # mypy struggles with torch.Size as a tuple of ints; cast explicitly.
+        return tuple(int(s) for s in self.probs.shape[:-1])
+
+    @property
+    def code_dim(self) -> int:
+        """Size of the final code dimension ``K``."""
+
+        return int(self.probs.shape[-1])

nervecode/integration/__init__.py

@@ -0,0 +1,9 @@
+"""Integration points and adapters.
+
+This subpackage will provide adapters and convenience utilities for integrating
+Nervecode into applications, scripts, and logging pipelines.
+"""
+
+from __future__ import annotations
+
+__all__: list[str] = []

nervecode/layers/__init__.py

@@ -0,0 +1,15 @@
+"""Layer wrappers and instrumentation shims.
+
+This subpackage will contain observe-only wrappers around selected PyTorch
+layers and utilities to apply them to user models.
+"""
+
+from __future__ import annotations
+
+from .base import BaseCodingWrapper, ReducerFn, identity_reduction
+
+__all__: list[str] = [
+    "BaseCodingWrapper",
+    "ReducerFn",
+    "identity_reduction",
+]

nervecode/layers/base.py

@@ -0,0 +1,333 @@
+"""Common wrapper base: bypass, trace cache, reduction, diagnostics hooks.
+
+This module defines a small, torch-agnostic base class and companion typing to
+standardize behavior shared by coding wrappers:
+
+- Bypass and fail-open controls (enable_coding(), disable_coding(), and the
+  bypass() context manager) with correct nesting semantics.
+- Latest per-layer trace caching (last_trace) with explicit update/clear.
+- A light reduction callable type compatible with identity or learned reducers.
+- Diagnostics hook registration that runs safely on trace updates.
+
+Wrappers built on top of this class remain observe-only: the base makes no
+assumptions about forward semantics beyond providing helpers that wrappers can
+use after computing the original layer output.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager, suppress
+from typing import Any
+
+try:  # Optional torch import for typing only; keep runtime torch-agnostic.
+    from nervecode.core import CodingTrace  # re-exported type
+except Exception:  # pragma: no cover - available in normal package use
+    CodingTrace = Any  # type: ignore[misc,assignment]
+
+__all__ = ["BaseCodingWrapper", "ReducerFn", "identity_reduction"]
+
+
+ReducerFn = Callable[[Any], tuple[Any, dict[str, Any]]]
+"""Callable that reduces an activation and returns metadata.
+
+Accepts a single activation-like object (typically a tensor with shape
+``(..., D_out)``) and returns ``(reduced, reduction_meta)`` where ``reduced``
+has shape ``(..., D)`` and ``reduction_meta`` is a JSON-serializable dict that
+identifies the applied reduction (e.g., {"method": "identity"}).
+"""
+
+
+def identity_reduction(x: Any) -> tuple[Any, dict[str, Any]]:
+    """Return input unchanged with minimal metadata."""
+
+    return x, {"method": "identity"}
+
+
+class BaseCodingWrapper:
+    """Torch-agnostic base for coding wrappers.
+
+    Responsibilities
+    - Maintain an enable/disable switch and a nestable bypass context to
+      implement fail-open behavior consistently across wrappers.
+    - Cache the most recent per-layer trace and run registered diagnostics
+      hooks safely when it updates.
+    - Provide a reducer slot with an identity default so wrappers can apply
+      a dimension reduction before coding without duplicating boilerplate.
+    """
+
+    def __init__(self, *, reducer: ReducerFn | None = None) -> None:
+        self._coding_enabled: bool = True
+        self._bypass_depth: int = 0
+        self._last_trace: CodingTrace | None = None
+        # Default to identity reduction when none is provided
+        self._reducer: ReducerFn = reducer or identity_reduction
+        self._diagnostic_hooks: list[Callable[[CodingTrace], None]] = []
+
+    # --- Bypass and enable/disable -------------------------------------------------
+    @property
+    def coding_enabled(self) -> bool:
+        """Return whether coding is globally enabled for this wrapper."""
+
+        return self._coding_enabled
+
+    @property
+    def is_coding_active(self) -> bool:
+        """True when coding is enabled and not currently bypassed."""
+
+        return self._coding_enabled and self._bypass_depth == 0
+
+    def enable_coding(self) -> None:
+        """Enable coding instrumentation for subsequent forwards."""
+
+        self._coding_enabled = True
+
+    def disable_coding(self) -> None:
+        """Disable coding instrumentation (fail-open) and clear caches.
+
+        Disabling coding should restore base-model behavior immediately. To
+        avoid stale reads from convenience accessors, drop any cached latest
+        trace at the time of disabling rather than waiting for the next
+        forward pass.
+        """
+
+        self._coding_enabled = False
+        # Clear any previously cached trace so that model-level aggregations
+        # observe the disabled state immediately.
+        self.clear_last_trace()
+
+    @contextmanager
+    def bypass(self) -> Iterator[None]:
+        """Temporarily bypass coding instrumentation.
+
+        The context is nestable; coding is considered bypassed when the bypass
+        depth is greater than zero.
+        """
+
+        self._bypass_depth += 1
+        try:
+            yield
+        finally:
+            self._bypass_depth -= 1
+            if self._bypass_depth < 0:  # pragma: no cover - defensive
+                self._bypass_depth = 0
+
+    # --- Reduction ----------------------------------------------------------------
+    def set_reducer(self, reducer: ReducerFn) -> None:
+        """Set the reducer callable used by ``maybe_reduce``."""
+
+        self._reducer = reducer
+
+    def get_reducer(self) -> ReducerFn:
+        """Return the current reducer callable."""
+
+        return self._reducer
+
+    def maybe_reduce(self, activation: Any) -> tuple[Any, dict[str, Any]]:
+        """Apply the configured reducer when coding is active.
+
+        When coding is not active (disabled or bypassed), return the input
+        unchanged with a metadata dictionary indicating bypass. This function
+        never raises; wrappers should keep observe-only semantics.
+        """
+
+        if not self.is_coding_active:
+            return activation, {"method": "bypass"}
+        try:
+            return self._reducer(activation)
+        except Exception:  # Fail-open on reducer errors
+            # Leave a recognizable marker to help downstream logging while
+            # guaranteeing observe-only behavior.
+            return activation, {"method": "bypass", "reason": "reduction_error"}
+
+    # --- Trace cache and diagnostics ----------------------------------------------
+    @property
+    def last_trace(self) -> CodingTrace | None:
+        """Return the most recently cached trace, if any."""
+
+        return self._last_trace
+
+    def clear_last_trace(self) -> None:
+        """Drop the cached trace."""
+
+        self._last_trace = None
+
+    def update_last_trace(self, trace: CodingTrace) -> None:
+        """Cache the latest trace and run diagnostics hooks safely."""
+
+        self._last_trace = trace
+        self._run_diagnostic_hooks(trace)
+
+    def add_diagnostics_hook(self, hook: Callable[[CodingTrace], None]) -> None:
+        """Register a hook invoked on every trace update.
+
+        Hooks must be side-effect safe; exceptions are suppressed to preserve
+        fail-open behavior.
+        """
+
+        self._diagnostic_hooks.append(hook)
+
+    def remove_diagnostics_hook(self, hook: Callable[[CodingTrace], None]) -> None:
+        """Remove a previously registered diagnostics hook, if present."""
+
+        with suppress(ValueError):
+            self._diagnostic_hooks.remove(hook)
+
+    def _run_diagnostic_hooks(self, trace: CodingTrace) -> None:
+        for hook in tuple(self._diagnostic_hooks):  # work on a snapshot
+            try:
+                hook(trace)
+            except Exception:
+                # Suppress hook exceptions to preserve user forwards; wrappers
+                # are observe-only. Hooks are best-effort diagnostics.
+                continue
+
+    # --- Layer-level diagnostics (latest-trace) ---------------------------------
+    def utilization(self) -> float | None:
+        """Fraction of codes used in the latest trace.
+
+        Returns the number of unique chosen indices divided by ``K`` from the
+        most recent trace, or ``None`` when unavailable. The result is a
+        Python ``float`` in ``[0, 1]``.
+        """
+
+        trace = self.last_trace
+        if trace is None:  # no trace cached
+            return None
+        try:
+            K = int(trace.soft_code.code_dim)
+            if K <= 0:
+                return None
+            # Prefer SoftCode.best_indices when present; otherwise fall back to
+            # the hard choice stored on the trace.
+            best = trace.soft_code.best_indices
+            if best is None:
+                best = trace.chosen_center_indices
+            values = self._flatten_to_list(best)
+            if values is None or K == 0:
+                return None
+            used = len({int(v) for v in values})
+            return float(used) / float(K)
+        except Exception:
+            return None
+
+    def mean_entropy(self) -> float | None:
+        """Mean assignment entropy from the latest trace, or ``None``.
+
+        Returns ``soft_code.entropy.mean()`` as a Python ``float`` when a
+        cached trace is available and contains entropy; otherwise ``None``.
+        """
+
+        trace = self.last_trace
+        if trace is None:
+            return None
+        entropy = trace.soft_code.entropy
+        if entropy is None:
+            return None
+        return self._mean_as_float(entropy)
+
+    def mean_code_length(self) -> float | None:
+        """Mean best-code length from the latest trace, or ``None``.
+
+        Returns ``soft_code.best_length.mean()`` as a Python ``float`` when a
+        cached trace is available and contains best-code lengths; otherwise
+        ``None``.
+        """
+
+        trace = self.last_trace
+        if trace is None:
+            return None
+        best_len = trace.soft_code.best_length
+        if best_len is None:
+            return None
+        return self._mean_as_float(best_len)
+
+    def mean_commitment_distance(self) -> float | None:
+        """Mean commitment distance from the latest trace, or ``None``.
+
+        Uses the per-position ``commitment_distances`` stored on the trace.
+        Returns a Python ``float`` when available, else ``None``.
+        """
+
+        trace = self.last_trace
+        if trace is None:
+            return None
+        d = trace.commitment_distances
+        if d is None:
+            return None
+        return self._mean_as_float(d)
+
+    @staticmethod
+    def _flatten_to_list(x: Any) -> list[float] | None:
+        """Best-effort flatten to a Python list of numbers.
+
+        Supports torch-like objects exposing ``reshape``/``view`` and ``tolist``,
+        as well as nested Python sequences.
+        """
+
+        # Torch-like path
+        try:
+            if hasattr(x, "reshape") and hasattr(x, "tolist"):
+                return list(x.reshape(-1).tolist())
+            if hasattr(x, "view") and hasattr(x, "tolist"):
+                return list(x.view(-1).tolist())
+            if hasattr(x, "tolist"):
+                out = x.tolist()
+                if isinstance(out, list):
+                    # Shallow list; attempt to flatten one level if needed
+                    def _flatten(v: Any) -> list[float]:
+                        if isinstance(v, list):
+                            res: list[float] = []
+                            for item in v:
+                                res.extend(_flatten(item))
+                            return res
+                        try:
+                            return [float(v)]
+                        except Exception:
+                            return []
+
+                    return _flatten(out)
+        except Exception:
+            pass
+
+        # Python sequence fallback
+        if isinstance(x, (list, tuple)):
+            res2: list[float] = []
+            for item in x:
+                sub = BaseCodingWrapper._flatten_to_list(item)
+                if sub is not None:
+                    res2.extend(sub)
+            return res2
+        # Scalar fallback
+        try:
+            return [float(x)]
+        except Exception:
+            return None
+
+    @staticmethod
+    def _mean_as_float(x: Any) -> float | None:
+        """Return mean(x) as a float for torch-like or Python sequences.
+
+        - If ``x`` has a ``mean()`` method, call it and extract ``item()`` if
+          available.
+        - Otherwise, attempt to flatten to a list and compute an arithmetic
+          mean.
+        """
+
+        # Torch-like: x.mean().item()
+        try:
+            if hasattr(x, "mean"):
+                m = x.mean()
+                if hasattr(m, "item"):
+                    return float(m.item())
+                return float(m)
+        except Exception:
+            pass
+
+        values = BaseCodingWrapper._flatten_to_list(x)
+        if not values:
+            return None
+        try:
+            return float(sum(values) / float(len(values)))
+        except Exception:
+            return None

nervecode/layers/conv.py

@@ -0,0 +1,174 @@
+"""Coding wrapper for ``torch.nn.Conv2d`` with pooled coding.
+
+This module provides ``CodingConv2d`` — an observe-only wrapper around
+``nn.Conv2d`` that computes coding statistics over a pooled representation of
+the layer output. The visible forward value remains identical to the underlying
+convolution; coding runs side-by-side and updates an internal trace cache.
+
+Design constraints for the first convolutional wrapper:
+- Observe-only contract for forwards (fail-open behavior).
+- Coding over a pooled representation (global average pooling over H and W).
+- Default sample-level reducer corresponds to global average pooling semantics.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .base import BaseCodingWrapper, ReducerFn
+
+# Import torch lazily but tolerate environments without it at import time.
+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)
+
+from nervecode.core import CodingTrace
+from nervecode.core.assignment import SoftAssignment
+from nervecode.core.codebook import Codebook
+
+__all__ = ["CodingConv2d"]
+
+
+class CodingConv2d(nn.Module, BaseCodingWrapper):
+    """Observe-only wrapper around ``nn.Conv2d`` with pooled coding.
+
+    The wrapper delegates computation to the underlying ``nn.Conv2d`` and,
+    when coding is active, computes a ``CodingTrace`` over a globally pooled
+    representation of the visible activation. The visible forward output is
+    never altered.
+
+    Parameters
+    - layer: The ``nn.Conv2d`` module to wrap.
+    - K: Number of codebook centers. Defaults to 16.
+    - coding_dim: Target coding dimension ``D``. When ``None`` (default), uses
+      ``layer.out_channels``. If provided, must satisfy ``1 <= coding_dim <=
+      layer.out_channels``; when smaller than ``out_channels``, a learned linear
+      projection is applied after global average pooling, affecting only the
+      coding path (observe-only contract preserved).
+    - reducer: Optional reducer callable. Overrides the default global-average
+      pooling behavior when provided.
+    - temperature: Soft-assignment temperature. Higher is softer.
+    - eps: Numerical epsilon to clamp divisions and logs.
+    """
+
+    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("CodingConv2d requires PyTorch to be installed")
+
+        nn.Module.__init__(self)
+        BaseCodingWrapper.__init__(self, reducer=reducer)
+
+        if not isinstance(layer, nn.Conv2d):
+            raise TypeError("layer must be an instance of torch.nn.Conv2d")
+
+        self.layer: Any = layer
+
+        out_channels = int(self.layer.out_channels)
+        if coding_dim is None:
+            code_dim = out_channels
+        else:
+            code_dim = int(coding_dim)
+            if code_dim <= 0:
+                raise ValueError("coding_dim must be >= 1")
+            if code_dim > out_channels:
+                raise ValueError("coding_dim must be <= layer.out_channels for reduction")
+
+        # Optional learned projection after pooling when reducing channels.
+        self._projection: Any | None = None
+        if reducer is None and code_dim < out_channels:
+            # Use a linear projection over the pooled channel vector (B, C) -> (B, D).
+            self._projection = nn.Linear(out_channels, code_dim, bias=False)
+
+            def _pool_then_project(x: Any) -> tuple[Any, dict[str, Any]]:
+                # Global average pooling over spatial dims (H, W)
+                pooled = x.mean(dim=(-1, -2))  # (B, C)
+                y_proj = self._projection(pooled)  # type: ignore[misc]
+                meta = {
+                    "method": "global_avg_pool2d+linear_projection",
+                    "in": out_channels,
+                    "out": code_dim,
+                    # Record that coding view comes from spatial reduction
+                    "spatial_reduction": True,
+                    "reduction_axes": [-2, -1],  # H, W in NCHW
+                }
+                return y_proj, meta
+
+            self.set_reducer(_pool_then_project)
+        elif reducer is None:
+            # Default: global average pooling over spatial dims to (B, C)
+            def _global_avg_pool2d(x: Any) -> tuple[Any, dict[str, Any]]:
+                reduced = x.mean(dim=(-1, -2))  # (B, C)
+                return reduced, {
+                    "method": "global_avg_pool2d",
+                    "spatial_reduction": True,
+                    "reduction_axes": [-2, -1],  # H, W in NCHW
+                }
+
+            self.set_reducer(_global_avg_pool2d)
+        # else: keep user-provided reducer
+
+        # Codebook over final feature dimension D=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.
+
+        Preserves the observe-only contract: the visible value is the
+        underlying layer output. For programmatic use, prefer
+        ``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."""
+
+        y = self.layer(x)
+
+        if not self.is_coding_active:
+            self.clear_last_trace()
+            return y, None
+
+        # Apply configured reducer (default: global average pooling over H, W).
+        reduced, reduction_meta = self.maybe_reduce(y)
+
+        # Compute assignments over the pooled representation.
+        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
+        status = "on" if self.coding_enabled else "off"
+        return (
+            f"layer=Conv2d(out_channels={int(self.layer.out_channels)}), "
+            f"code_dim={int(self.codebook.code_dim)}, "
+            f"K={int(self.codebook.K)}, reducer={self.get_reducer().__name__}, "
+            f"coding={status}"
+        )