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

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
import json
import math
from pathlib import Path
from typing import Any

import torch
from torch import nn


@dataclass(frozen=True)
class AxiomDimConfig:
    """Procedural virtual dimension for tiny native models.

    ``effective_dim`` is the symbolic reasoning/coordinate capacity. The bridge
    only materializes ``physical_dim`` activations and ``basis_rank`` latent
    coordinates, so it can target very wide dimensions without dense blowup.
    """

    physical_dim: int = 128
    effective_dim: int = 20_480
    basis_rank: int = 64
    facets: int = 8
    residual_scale: float = 0.25


class AxiomDimBridge(nn.Module):
    """TinyMind AxiomDim: procedural high-dimensional capacity in small tensors.

    The key trick is to model a huge dimension as reusable rank coordinates plus
    facet gates. No [batch, seq, effective_dim] tensor is ever created.
    """

    def __init__(self, cfg: AxiomDimConfig):
        super().__init__()
        if cfg.physical_dim <= 0 or cfg.effective_dim <= 0 or cfg.basis_rank <= 0 or cfg.facets <= 0:
            raise ValueError("AxiomDim dimensions must be positive")
        self.cfg = cfg
        self.to_basis = nn.Linear(cfg.physical_dim, cfg.basis_rank, bias=False)
        self.from_basis = nn.Linear(cfg.basis_rank, cfg.physical_dim, bias=False)
        self.facet_gate = nn.Parameter(torch.zeros(cfg.facets, cfg.basis_rank))
        phase = torch.linspace(0.0, math.pi, cfg.basis_rank)
        self.register_buffer("procedural_phase", phase, persistent=False)

    @property
    def parameter_count(self) -> int:
        return sum(param.numel() for param in self.parameters())

    def dense_dim_params_estimate(self) -> int:
        return 2 * self.cfg.physical_dim * self.cfg.effective_dim

    def compression_vs_dense_dim(self) -> float:
        return self.dense_dim_params_estimate() / max(1, self.parameter_count)

    def forward(self, x: torch.Tensor, *, return_state: bool = False):
        z = self.to_basis(x)
        phase = self.procedural_phase.to(dtype=z.dtype, device=z.device)
        facet = torch.sigmoid(self.facet_gate).mean(dim=0).to(dtype=z.dtype, device=z.device)
        latent = torch.tanh((z + torch.sin(phase)) * facet)
        y = x + self.cfg.residual_scale * self.from_basis(latent)
        if not return_state:
            return y
        state = {
            "effective_dim": self.cfg.effective_dim,
            "physical_dim": self.cfg.physical_dim,
            "basis_rank": self.cfg.basis_rank,
            "facets": self.cfg.facets,
            "materializes_effective_dim": False,
            "latent_shape": list(latent.shape),
            "parameter_count": self.parameter_count,
            "compression_vs_dense_dim": self.compression_vs_dense_dim(),
        }
        return y, state


def _smoke(cfg: AxiomDimConfig) -> dict[str, Any]:
    torch.manual_seed(20260527)
    bridge = AxiomDimBridge(cfg)
    x = torch.randn(2, 8, cfg.physical_dim, requires_grad=True)
    y, state = bridge(x, return_state=True)
    loss = y.float().pow(2).mean()
    loss.backward()
    grads = [param.grad for param in bridge.parameters() if param.grad is not None]
    backward_finite = bool(grads) and all(torch.isfinite(grad).all().item() for grad in grads)
    backward_finite = backward_finite and x.grad is not None and bool(torch.isfinite(x.grad).all().item())
    return {
        "forward_finite": bool(torch.isfinite(y).all().item()),
        "backward_finite": backward_finite,
        "loss": float(loss.detach().cpu()),
        "state": state,
    }


def _candidate(effective_dim: int, physical_dim: int, basis_rank: int, facets: int) -> dict[str, Any]:
    cfg = AxiomDimConfig(
        physical_dim=physical_dim,
        effective_dim=effective_dim,
        basis_rank=basis_rank,
        facets=facets,
    )
    bridge = AxiomDimBridge(cfg)
    smoke = _smoke(cfg)
    compression = bridge.compression_vs_dense_dim()
    return {
        "effective_dim": effective_dim,
        "physical_dim": physical_dim,
        "basis_rank": basis_rank,
        "facets": facets,
        "materializes_effective_dim": False,
        "parameter_count": bridge.parameter_count,
        "dense_dim_params_estimate": bridge.dense_dim_params_estimate(),
        "compression_vs_dense_dim": compression,
        "smoke": smoke,
        "score": math.log1p(compression) + math.log1p(physical_dim) + 0.5 * math.log1p(basis_rank) + 0.25 * math.log1p(facets),
    }


def build_axiomdim_report(
    out_dir: str | Path,
    *,
    effective_dim: int = 20_480,
    physical_dims: list[int] | None = None,
    basis_ranks: list[int] | None = None,
    facets: list[int] | None = None,
) -> dict[str, Any]:
    physical_values = physical_dims or [128, 256, 512]
    rank_values = basis_ranks or [32, 64, 96, 128]
    facet_values = facets or [8, 16, 32]
    candidates = [
        _candidate(effective_dim, physical_dim, rank, facet)
        for physical_dim in physical_values
        for rank in rank_values
        for facet in facet_values
    ]
    best = max(candidates, key=lambda item: item["score"])
    report = {
        "schema": "tinymind.axiomdim.v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "target": {
            "effective_dim": effective_dim,
            "method": "AxiomDim procedural rank-facet dimension",
            "materializes_effective_dim": False,
        },
        "summary": {
            "candidate_count": len(candidates),
            "physical_dims": physical_values,
            "basis_ranks": rank_values,
            "facets": facet_values,
        },
        "best_candidate": best,
        "top_candidates": sorted(candidates, key=lambda item: item["score"], reverse=True)[:8],
        "claim_gate": {
            "axiomdim_candidate_ready": best["smoke"]["forward_finite"] and best["smoke"]["backward_finite"],
            "dense_dim_claim_allowed": False,
            "tier0_claim_allowed": False,
            "world_best_claim_allowed": False,
            "reason": "AxiomDim proves a factorized procedural dimension candidate, not an externally validated frontier model.",
        },
    }
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    path = out / "axiomdim_report.json"
    report["json_path"] = str(path)
    path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return report