tinymind-native-colab-handoff/bundle/evaluation/layer_coherence.py

"""Layer coherence auditor for complex TinyMind architectures."""

from __future__ import annotations

from datetime import datetime, timezone
import json
from pathlib import Path

import torch

from model.axiom_weave import AxiomWeaveModel
from model.config import axiomweave_config


def audit_axiomweave_coherence(out_dir: str | Path, size: str = "tiny", seq_len: int = 16) -> dict:
    cfg = axiomweave_config(size)
    model = AxiomWeaveModel(cfg)
    ids = torch.randint(4, min(cfg.vocab_size, 4096), (1, int(seq_len)))
    with torch.no_grad():
        out = model(ids, return_stats=True)
    rows = []
    dead_layers = []
    route_floor = 1e-5
    norm_floor = 1e-8
    for idx, stats in enumerate(out["stats"]):
        route = [float(x) for x in stats["route_weights_mean"].cpu()]
        norms = [float(x) for x in stats["branch_norms"].cpu()]
        memory_norm = float(stats["purefield_memory_norm"].cpu())
        route_alive = all(value > route_floor for value in route)
        branch_alive = all(value > norm_floor for value in norms)
        finite = all(torch.isfinite(t).all().item() for t in [stats["route_weights_mean"], stats["branch_norms"], stats["purefield_memory_norm"]])
        alive = route_alive and branch_alive and finite
        if not alive:
            dead_layers.append(idx)
        rows.append(
            {
                "layer": idx,
                "route_attention": route[0],
                "route_ssm": route[1],
                "route_purefield": route[2],
                "branch_norm_attention": norms[0],
                "branch_norm_ssm": norms[1],
                "branch_norm_purefield": norms[2],
                "purefield_memory_norm": memory_norm,
                "route_entropy": float(stats["route_entropy"].cpu()),
                "alive": alive,
            }
        )
    finite_logits = bool(torch.isfinite(out["logits"]).all().item())
    alive_ratio = sum(1 for row in rows if row["alive"]) / max(len(rows), 1)
    entropy_values = [row["route_entropy"] for row in rows]
    entropy_mean = sum(entropy_values) / max(len(entropy_values), 1)
    harmony_score = 100.0 * alive_ratio * min(1.0, entropy_mean / 0.5) if finite_logits else 0.0
    report = {
        "schema_version": "tinymind-layer-coherence-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "architecture": "AxiomWeave",
        "size": size,
        "seq_len": int(seq_len),
        "finite_logits": finite_logits,
        "layers": rows,
        "dead_layers": dead_layers,
        "no_zero_work_gate": {
            "passed": finite_logits and not dead_layers,
            "alive_layer_ratio": alive_ratio,
            "route_floor": route_floor,
            "branch_norm_floor": norm_floor,
        },
        "harmony_score": harmony_score,
        "purity_note": "Every layer is audited for nonzero branch energy and finite coordinated routing; no world-best claim is made.",
        "world_best_claim_allowed": False,
    }
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    json_path = out / "layer_coherence_report.json"
    md_path = out / "layer_coherence_report.md"
    report["json_path"] = str(json_path)
    report["markdown_path"] = str(md_path)
    json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    md_path.write_text(_markdown(report), encoding="utf-8")
    return report


def _markdown(report: dict) -> str:
    return "\n".join(
        [
            "# TinyMind Layer Coherence Report",
            "",
            f"- Architecture: {report['architecture']}",
            f"- No-zero-work gate: {report['no_zero_work_gate']['passed']}",
            f"- Alive layer ratio: {report['no_zero_work_gate']['alive_layer_ratio']:.2%}",
            f"- Harmony score: {report['harmony_score']:.2f}",
            f"- Dead layers: {report['dead_layers']}",
            "- World-best claim: false",
            "",
        ]
    )