example.py

"""Minimal end-to-end example: build, forward, generate, and print the LTI
spectral radius bound for both GQA and MLA attention variants.
"""

import torch

from open_mythos.main import MythosConfig, OpenMythos


def build(attn_type: str) -> MythosConfig:
    base = dict(
        vocab_size=1000,
        dim=256,
        n_heads=8,
        max_seq_len=128,
        max_loop_iters=4,
        prelude_layers=1,
        coda_layers=1,
        n_experts=8,
        n_shared_experts=1,
        n_experts_per_tok=2,
        expert_dim=64,
        lora_rank=8,
        attn_type=attn_type,
    )
    if attn_type == "gqa":
        return MythosConfig(**base, n_kv_heads=2)
    return MythosConfig(
        **base,
        n_kv_heads=8,
        kv_lora_rank=32,
        q_lora_rank=64,
        qk_rope_head_dim=16,
        qk_nope_head_dim=16,
        v_head_dim=16,
    )


def demo(attn_type: str) -> None:
    torch.manual_seed(0)
    cfg = build(attn_type)
    model = OpenMythos(cfg)

    ids = torch.randint(0, cfg.vocab_size, (2, 16))
    logits = model(ids, n_loops=4)
    out = model.generate(ids, max_new_tokens=8, n_loops=8)
    A = model.recurrent.injection.get_A()

    tag = attn_type.upper()
    print(f"\n[{tag}] Parameters:   {sum(p.numel() for p in model.parameters()):,}")
    print(f"[{tag}] Logits shape: {tuple(logits.shape)}")
    print(f"[{tag}] Generated:    {tuple(out.shape)}")
    print(f"[{tag}] \u03c1(A) max:     {A.max().item():.4f}  (guaranteed <= 1)")
    print(f"[{tag}] Ponder mean:  {model.recurrent.last_ponder_cost.mean().item():.2f}")


if __name__ == "__main__":
    for attn in ("gqa", "mla"):
        demo(attn)