tinymind-native-8b-remote-handoff/bundle/model/pure_lattice_cnn.py

"""PureLattice CNN core for TinyMind.

The core is a compact convolutional reasoning adapter over hidden states. It is
designed as a measurable plug-in, not as a claim that CNN alone replaces the
language model. Token streams and feature grids are projected through the same
multi-scale lattice so local structure, code syntax, and extracted multimodal
features can be sharpened before the recurrent/attention stack consumes them.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

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


@dataclass(frozen=True)
class PureLatticeCNNConfig:
    dim: int
    hidden_mult: int = 2
    kernel_sizes: tuple[int, ...] = (3, 5, 9)
    dilations: tuple[int, ...] = (1, 2, 4)
    dropout: float = 0.0
    residual_scale: float = 0.5
    eps: float = 1e-5

    def __post_init__(self) -> None:
        if self.dim <= 0:
            raise ValueError("dim must be positive")
        if self.hidden_mult <= 0:
            raise ValueError("hidden_mult must be positive")
        if len(self.kernel_sizes) != len(self.dilations):
            raise ValueError("kernel_sizes and dilations must have the same length")
        if not self.kernel_sizes:
            raise ValueError("at least one convolution branch is required")
        if any(k < 3 or k % 2 == 0 for k in self.kernel_sizes):
            raise ValueError("kernel sizes must be odd integers >= 3")
        if any(d <= 0 for d in self.dilations):
            raise ValueError("dilations must be positive")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must be in [0, 1)")
        if not 0.0 < self.residual_scale <= 1.0:
            raise ValueError("residual_scale must be in (0, 1]")


class _DepthwiseBranch(nn.Module):
    def __init__(self, dim: int, kernel_size: int, dilation: int):
        super().__init__()
        padding = dilation * (kernel_size - 1) // 2
        self.conv = nn.Conv1d(dim, dim, kernel_size, padding=padding, dilation=dilation, groups=dim, bias=False)
        self.norm = nn.GroupNorm(1, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.norm(self.conv(x))


class PureLatticeCNNCore(nn.Module):
    """Multi-scale gated CNN adapter for token and feature-grid hidden states."""

    def __init__(self, cfg: PureLatticeCNNConfig):
        super().__init__()
        self.cfg = cfg
        self.in_norm = nn.LayerNorm(cfg.dim, eps=cfg.eps)
        self.branches = nn.ModuleList(
            _DepthwiseBranch(cfg.dim, kernel, dilation)
            for kernel, dilation in zip(cfg.kernel_sizes, cfg.dilations)
        )
        branch_dim = cfg.dim * len(cfg.kernel_sizes)
        hidden_dim = cfg.dim * cfg.hidden_mult
        self.mix = nn.Sequential(
            nn.Linear(branch_dim, hidden_dim, bias=False),
            nn.SiLU(),
            nn.Dropout(cfg.dropout),
            nn.Linear(hidden_dim, cfg.dim, bias=False),
        )
        self.gate = nn.Linear(cfg.dim, cfg.dim, bias=True)
        self.out_norm = nn.LayerNorm(cfg.dim, eps=cfg.eps)
        self._reset_parameters()

    @property
    def receptive_field(self) -> int:
        return 1 + sum((kernel - 1) * dilation for kernel, dilation in zip(self.cfg.kernel_sizes, self.cfg.dilations))

    def forward(self, hidden: torch.Tensor) -> tuple[torch.Tensor, dict[str, Any]]:
        if hidden.ndim != 3:
            raise ValueError("hidden must be shaped [batch, tokens, dim]")
        if hidden.shape[-1] != self.cfg.dim:
            raise ValueError(f"expected dim {self.cfg.dim}, got {hidden.shape[-1]}")

        u = self.in_norm(hidden)
        conv_in = u.transpose(1, 2).contiguous()
        branch_outputs = [branch(conv_in).transpose(1, 2).contiguous() for branch in self.branches]
        mixed = self.mix(torch.cat(branch_outputs, dim=-1))
        gate = torch.sigmoid(self.gate(u))
        delta = torch.tanh(mixed) * gate
        out = self.out_norm(hidden + self.cfg.residual_scale * delta)
        state = self._state(out, gate, input_kind="tokens")
        return out, state

    def forward_grid(self, grid: torch.Tensor) -> tuple[torch.Tensor, dict[str, Any]]:
        if grid.ndim != 4:
            raise ValueError("grid must be shaped [batch, channels, height, width]")
        batch, channels, height, width = grid.shape
        if channels != self.cfg.dim:
            raise ValueError(f"expected channels {self.cfg.dim}, got {channels}")
        tokens = grid.flatten(2).transpose(1, 2).contiguous()
        out_tokens, state = self.forward(tokens)
        out_grid = out_tokens.transpose(1, 2).reshape(batch, channels, height, width).contiguous()
        state = {
            **state,
            "input_kind": "grid_2d",
            "height": height,
            "width": width,
            "native_vision_claim_allowed": False,
        }
        return out_grid, state

    def _state(self, out: torch.Tensor, gate: torch.Tensor, input_kind: str) -> dict[str, Any]:
        return {
            "input_kind": input_kind,
            "branch_count": len(self.branches),
            "receptive_field": self.receptive_field,
            "output_energy": float(out.detach().float().pow(2).mean().item()),
            "gate_min": float(gate.detach().min().item()),
            "gate_max": float(gate.detach().max().item()),
            "cnn_core_ready": bool(torch.isfinite(out).all().item()),
            "world_best_cnn_claim_allowed": False,
        }

    def _reset_parameters(self) -> None:
        for module in self.modules():
            if isinstance(module, (nn.Linear, nn.Conv1d)):
                nn.init.xavier_uniform_(module.weight)
                if getattr(module, "bias", None) is not None:
                    nn.init.zeros_(module.bias)
        nn.init.constant_(self.gate.bias, -1.0)


def count_trainable_parameters(module: nn.Module) -> int:
    return sum(param.numel() for param in module.parameters() if param.requires_grad)