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

"""Local evidence benchmark for AxiomFlow HyperWeave."""

from __future__ import annotations

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

import torch

from model.axiomflow import AxiomFlowBlock, AxiomFlowConfig


def _jsonify_pressure(pressure: dict) -> dict:
    return {
        "aux_loss": float(pressure["aux_loss"].detach().cpu()),
        "route_balance_loss": float(pressure["route_balance_loss"].cpu()),
        "lane_diversity_loss": float(pressure["lane_diversity_loss"].cpu()),
        "contribution_loss": float(pressure["contribution_loss"].cpu()),
        "anti_collapse_gate": pressure["anti_collapse_gate"],
        "lane_diversity_gate": pressure["lane_diversity_gate"],
    }


def run_axiomflow_bench(
    out_dir: str | Path,
    dim: int = 128,
    seq_len: int = 64,
    batch_size: int = 1,
    local_window: int = 16,
    memory_slots: int = 8,
    memory_rank: int = 16,
    retrieval_top_k: int = 4,
    retrieved_chunk_tokens: int = 32,
    seed: int = 20260525,
) -> dict:
    torch.manual_seed(int(seed))
    cfg = AxiomFlowConfig(
        dim=int(dim),
        local_window=int(local_window),
        memory_slots=int(memory_slots),
        memory_rank=int(memory_rank),
        retrieval_top_k=int(retrieval_top_k),
        retrieved_chunk_tokens=int(retrieved_chunk_tokens),
        residual_alpha=0.2,
    )
    block = AxiomFlowBlock(cfg)
    hidden = torch.randn(int(batch_size), int(seq_len), cfg.dim, requires_grad=True)
    retrieved = torch.randn(
        int(batch_size),
        cfg.retrieval_top_k,
        cfg.retrieved_chunk_tokens,
        cfg.dim,
    )
    start = time.perf_counter()
    out, state, metrics = block(hidden, retrieved_chunks=retrieved, return_metrics=True)
    forward_seconds = time.perf_counter() - start
    loss = out.float().pow(2).mean()
    loss.backward()
    grad_tensors = [p.grad for p in block.parameters() if p.grad is not None]
    backward_finite = bool(all(torch.isfinite(g).all().item() for g in grad_tensors))
    forward_finite = bool(torch.isfinite(out).all().item())
    full_kv_tokens = int(seq_len)
    bounded_tokens = int(state.cached_token_capacity())
    report = {
        "schema_version": "tinymind-axiomflow-bench-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "architecture": "AxiomFlow HyperWeave",
        "dim": cfg.dim,
        "seq_len": int(seq_len),
        "batch_size": int(batch_size),
        "forward_seconds": forward_seconds,
        "forward_finite": forward_finite,
        "backward_finite": backward_finite,
        "route_weights_mean": [float(x) for x in metrics["route_weights_mean"].cpu()],
        "route_entropy": float(metrics["route_entropy"].cpu()),
        "local_exact_tokens_stored": int(metrics["local_exact_tokens_stored"]),
        "long_context_kv_tokens_stored": int(metrics["long_context_kv_tokens_stored"]),
        "bounded_state_tokens_equivalent": int(metrics["bounded_state_tokens_equivalent"]),
        "full_kv_tokens_reference": full_kv_tokens,
        "bounded_memory_gate": {
            "passed": bool(bounded_tokens <= cfg.local_window + cfg.memory_slots),
            "cached_token_capacity": bounded_tokens,
            "max_cached_token_capacity": int(cfg.local_window + cfg.memory_slots),
            "long_context_kv_tokens_stored": 0,
        },
        "coherence_gate": metrics["coherence_gate"],
        "intensity_pressure": _jsonify_pressure(metrics["intensity_pressure"]),
        "claim_gate": {
            "new_reference_implementation": True,
            "beats_flashattention3_claim_allowed": False,
            "beats_deepspeed_ulysses_claim_allowed": False,
            "beats_xformers_claim_allowed": False,
            "beats_ringattention_claim_allowed": False,
            "reason": (
                "AxiomFlow v1 has local finite/bounded-memory evidence only. "
                "Superiority claims require controlled CUDA kernels and external benchmark comparisons."
            ),
        },
    }
    out_path = Path(out_dir)
    out_path.mkdir(parents=True, exist_ok=True)
    json_path = out_path / "axiomflow_bench_report.json"
    md_path = out_path / "axiomflow_bench_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:
    lines = [
        "# TinyMind AxiomFlow HyperWeave Bench",
        "",
        f"- Forward finite: {report['forward_finite']}",
        f"- Backward finite: {report['backward_finite']}",
        f"- Cached token capacity: {report['bounded_memory_gate']['cached_token_capacity']}",
        f"- Long-context KV tokens stored: {report['long_context_kv_tokens_stored']}",
        f"- Coherence gate: {report['coherence_gate']['passed']}",
        f"- Mutual support score: {report['coherence_gate']['mutual_support_score']:.4f}",
        f"- Intensity aux loss: {report['intensity_pressure']['aux_loss']:.6f}",
        f"- Anti-collapse gate: {report['intensity_pressure']['anti_collapse_gate']['passed']}",
        f"- Lane-diversity gate: {report['intensity_pressure']['lane_diversity_gate']['passed']}",
        "- Superiority claims: blocked until fair benchmark evidence exists",
        "",
        "## Route Weights",
        "",
        f"- local_exact: {report['route_weights_mean'][0]:.4f}",
        f"- ledger_recall: {report['route_weights_mean'][1]:.4f}",
        f"- compressed_field: {report['route_weights_mean'][2]:.4f}",
    ]
    return "\n".join(lines) + "\n"