px_modules.py

"""
gemma3-px  —  Pure Architectural Modules
==========================================
Zero-shot only. No trainable parameters beyond what is strictly necessary.
No OpenMythos baggage (no ACT, no LoRA, no probes, no epsilon).

Three modules:
  LTIInjection       — the core recurrence signal: h_new = trans_out + gamma*(e_norm - h)
  ADCInjection       — adaptive variant: gamma scales with instability (1-phi)
  StabilityMonitor   — parameter-free Phi / Lambda / Eta heuristics (logging only)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional


# ---------------------------------------------------------------------------
# 1. LTI Injection  (Pure, fixed gamma)
# ---------------------------------------------------------------------------

class LTIInjection(nn.Module):
    """
    Linear Time-Invariant anchor injection.
    Provides 'Computational Headroom' by pulling h back toward the
    frozen input embedding e whenever it drifts.

    Formula:
        h_new = transformer_out + gamma * (LayerNorm(e) - h)

    Identity at t=0 if gamma=0. Optimal empirical gamma: 0.08.
    """

    def __init__(self, dim: int, gamma: float = 0.08):
        super().__init__()
        self.gamma = gamma
        # Affine=False: no trainable params, pure normalization
        self.input_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)

    def forward(
        self,
        h: torch.Tensor,
        e: torch.Tensor,
        transformer_out: torch.Tensor,
    ) -> torch.Tensor:
        e_norm = self.input_norm(e.to(torch.float32)).to(h.dtype)
        return transformer_out + self.gamma * (e_norm - h)


# ---------------------------------------------------------------------------
# 2. ADC Injection  (Adaptive Dynamic Correction)
# ---------------------------------------------------------------------------

class ADCInjection(nn.Module):
    """
    Adaptive variant of LTI: effective gamma = base_gamma + alpha*(1-phi).

    When phi is high (state stable) → injection is gentle.
    When phi drops (state drifting) → injection strengthens automatically.
    No trainable parameters.

    Args:
        dim:        hidden size
        base_gamma: minimum injection strength (default 0.06)
        alpha:      maximum additional strength at full instability (default 0.10)
    """

    def __init__(self, dim: int, base_gamma: float = 0.06, alpha: float = 0.10):
        super().__init__()
        self.base_gamma = base_gamma
        self.alpha = alpha
        self.input_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)

    def forward(
        self,
        h: torch.Tensor,
        e: torch.Tensor,
        transformer_out: torch.Tensor,
        phi: float = 1.0,          # Φ from previous loop (1.0 = fully stable)
    ) -> torch.Tensor:
        instability = max(0.0, 1.0 - phi)
        effective_gamma = self.base_gamma + self.alpha * instability
        e_norm = self.input_norm(e.to(torch.float32)).to(h.dtype)
        return transformer_out + effective_gamma * (e_norm - h)


# ---------------------------------------------------------------------------
# 3. Stability Monitor  (no parameters — logging / diagnostics only)
# ---------------------------------------------------------------------------

class StabilityMonitor:
    """
    Parameter-free heuristics for Phi (Φ), Lambda (λ), and Eta (η).

    Phi   — cosine similarity between consecutive hidden states (stability).
    Lambda— cosine distance between current h and anchor e (drift).
    Eta   — variance-based entropy estimate (uncertainty).
    """

    @staticmethod
    def calculate_phi(h_new: torch.Tensor, h_old: torch.Tensor) -> torch.Tensor:
        """Φ ∈ [0,1]: 1 = identical state, 0 = orthogonal."""
        B = h_new.shape[0]
        # Safe FP32 for FP16 models
        h_n_f32 = h_new.view(B, -1).to(torch.float32)
        h_o_f32 = h_old.view(B, -1).to(torch.float32)
        return F.cosine_similarity(h_n_f32, h_o_f32, dim=-1).mean()

    @staticmethod
    def detect_lambda(h: torch.Tensor, e: torch.Tensor) -> torch.Tensor:
        """λ ∈ [0,1]: 0 = h == e (no drift), 1 = fully diverged."""
        B = h.shape[0]
        sim = F.cosine_similarity(h.view(B, -1), e.view(B, -1), dim=-1).mean()
        return (1.0 - sim).clamp(0.0, 1.0)

    @staticmethod
    def calculate_eta(h: torch.Tensor) -> torch.Tensor:
        """η ∈ [0,1]: variance-based uncertainty proxy. Recalibrated for Gemma-3 (var~67)."""
        var = torch.var(h.to(torch.float32), dim=-1).mean()
        # Scale 67.5 to ~0.5. 
        # (67.5 / 67.5) * 10 - 5 = 5 (sigmoid -> 0.99)
        # (50 / 67.5) * 10 - 5 = 2.4 (sigmoid -> 0.9)
        # (30 / 67.5) * 10 - 5 = -0.5 (sigmoid -> 0.37)
        return torch.sigmoid((var / 67.5) * 10.0 - 5.0)