00_README.md

Exact Muon Polar Express MPS optimization

This bundle contains the independent frozen reference, candidate kernels, exact evaluator, public API benchmark, safety audit, and final receipts for the Apple Silicon MPS optimization.

Run from the repository root with the project environment:

PYTHONPATH=. python optimization_candidates/eval_muon_exact.py --device mps --candidate out_mm_bx_cast_after_context
PYTHONPATH=. python optimization_candidates/eval_muon_public_api.py
PYTHONPATH=. python optimization_candidates/audit_muon_safety.py

The final public benchmark receipt reports a 1.091x median speedup. The gates use raw tensor-byte equality, output metadata, input immutability, contracts, gradients, adversarial values, and held-out shapes.

audit_muon_safety.py

"""Independent safety audit for the public Muon polar-express API.

This audit is intentionally separate from the performance evaluator.  It
checks the production route against the frozen reference on adversarial values
and tensor metadata, and statically rejects candidate/evaluator shortcuts that
could hide a baseline fallback or benchmark-fixture detection.
"""

from __future__ import annotations

import ast
import importlib.util
import json
import random
import sys
from pathlib import Path

import torch


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "optimization_candidates"))
from baseline_muon_frozen import polar_express_frozen

from optimizers.muon import zeropower_polar_express


def _load_candidates():
    path = ROOT / "optimization_candidates" / "muon_exact_candidates.py"
    spec = importlib.util.spec_from_file_location("muon_exact_candidates_audit", path)
    if spec is None or spec.loader is None:
        raise RuntimeError("candidate import failed")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return {
        "contiguous_tall_only": module.contiguous_tall_only,
        "cast_after_transpose_tall": module.cast_after_transpose_tall,
        "clone_final_tall": module.clone_final_tall,
        "out_mm_bx_only": module.out_mm_bx_only,
        "out_mm_a2_only": module.out_mm_a2_only,
        "out_mm_bx_muon_context": module.out_mm_bx_muon_context,
        "out_mm_bx_cast_after_context": module.out_mm_bx_cast_after_context,
    }


def _inputs(device):
    values = {
        "zeros": torch.zeros((1536, 512), device=device, dtype=torch.float32),
        "ones": torch.ones((1536, 512), device=device, dtype=torch.float32),
        "large": torch.full((1536, 512), 1e10, device=device, dtype=torch.float32),
        "small": torch.full((1536, 512), 1e-10, device=device, dtype=torch.float32),
        "mixed": torch.tensor([0.0, -0.0, float("inf"), float("-inf"), float("nan")], device=device, dtype=torch.float32).repeat(1536, 512 // 5 + 1)[:, :512],
    }
    base = torch.arange(1536 * 512, device=device, dtype=torch.float32).reshape(1536, 512)
    values["pattern"] = (base.remainder(17) - 8) / 3
    storage = torch.empty((1536, 1024), device=device, dtype=torch.float32)
    storage[..., ::2] = values["pattern"]
    values["noncontiguous"] = storage[..., ::2]
    return values


def _same(a, b):
    if a.shape != b.shape or a.dtype != b.dtype or a.device.type != b.device.type:
        return False
    # Compare actual bytes so NaN payloads and signed zero are both part of the
    # exact contract; numeric equality alone would accept changed bit patterns.
    try:
        return torch.equal(
            a.detach().cpu().clone(memory_format=torch.contiguous_format).view(torch.uint8),
            b.detach().cpu().clone(memory_format=torch.contiguous_format).view(torch.uint8),
        )
    except (RuntimeError, ValueError):
        return False


def _metadata(tensor: torch.Tensor, input_tensor: torch.Tensor) -> tuple:
    return (
        tuple(tensor.shape),
        tensor.dtype,
        tensor.device.type,
        tensor.device.index,
        tuple(tensor.stride()),
        tensor.storage_offset(),
        tensor.requires_grad,
        tensor.data_ptr() == input_tensor.data_ptr(),
    )


def _static_audit():
    source_paths = [ROOT / "optimization_candidates" / "muon_exact_candidates.py", ROOT / "optimization_candidates" / "eval_muon_exact.py", ROOT / "optimization_candidates" / "eval_muon_workload.py"]
    findings = []
    for path in source_paths:
        tree = ast.parse(path.read_text(), filename=str(path))
        names = {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)}
        # Candidate source may contain descriptive text, but not invoke the
        # mutable production baseline or call a hidden fallback.
        if path.name == "muon_exact_candidates.py" and "_zeropower_polar_express_impl" in names:
            findings.append(f"candidate references production baseline: {path}")
        if path.name == "muon_exact_candidates.py" and {"eval_muon_exact", "baseline_muon_frozen"} & names:
            findings.append(f"candidate imports evaluator/reference: {path}")
    return {"passed": not findings, "findings": findings}


def _contract_audit(public, candidate, device):
    x = torch.randn((1536, 512), device=device, dtype=torch.float32)
    checks = {}
    for name, fn in (("public", public), ("candidate", candidate)):
        invalid_step = False
        try:
            fn(x, steps=6)
        except (AssertionError, RuntimeError):
            invalid_step = True
        invalid_rank = False
        try:
            fn(torch.randn((8,), device=device), steps=5)
        except (AssertionError, RuntimeError):
            invalid_rank = True
        zero = fn(x, steps=0)
        negative = fn(x, steps=-1)
        checks[name] = {
            "invalid_steps_rejected": invalid_step,
            "invalid_rank_rejected": invalid_rank,
            "steps_zero_shape": list(zero.shape),
            "negative_steps_shape": list(negative.shape),
        }
    return checks


def _public_gradient_audit(public, device):
    """Check production autograd against the frozen reference independently."""
    cases = [((13, 7), 5), ((7, 13), 5), ((2, 7, 13), 3), ((32, 64), 1)]
    rows = []
    for shape, steps in cases:
        seed = 73000 + len(rows)
        generator = torch.Generator(device="cpu").manual_seed(seed)
        x0 = torch.randn(shape, generator=generator, dtype=torch.float32).to(device)
        weight = torch.randn(shape, generator=torch.Generator(device="cpu").manual_seed(seed + 1), dtype=torch.bfloat16).to(device)
        x_ref = x0.clone().requires_grad_(True)
        x_public = x0.clone().requires_grad_(True)
        try:
            ref = (polar_express_frozen(x_ref, steps=steps) * weight).sum()
            out = (public(x_public, steps=steps) * weight).sum()
            ref.backward()
            out.backward()
            gradients_exact = _same(x_ref.grad, x_public.grad)
            execution_error = None
        except (RuntimeError, AssertionError) as exc:
            gradients_exact = False
            execution_error = f"{type(exc).__name__}: {exc}"
        rows.append({
            "shape": list(shape),
            "steps": steps,
            "gradients_exact": gradients_exact,
            "execution_error": execution_error,
        })
    return {
        "passed": all(row["gradients_exact"] and row["execution_error"] is None for row in rows),
        "cases": rows,
    }


def main():
    if not torch.backends.mps.is_available():
        raise SystemExit("MPS unavailable")
    device = torch.device("mps")
    candidates = _load_candidates()
    rows = []
    for candidate_name, candidate in candidates.items():
      for name, x in _inputs(device).items():
        before = x.clone()
        ref = polar_express_frozen(x, steps=5)
        out = zeropower_polar_express(x, steps=5)
        candidate_out = candidate(x, steps=5)
        ref_meta = _metadata(ref, x)
        public_meta = _metadata(out, x)
        candidate_meta = _metadata(candidate_out, x)
        rows.append({
            "candidate": candidate_name,
            "case": name,
            "public_exact": _same(ref, out),
            "candidate_exact": _same(ref, candidate_out),
            "input_unchanged": _same(x, before),
            "public_shape": list(out.shape),
            "candidate_shape": list(candidate_out.shape),
            "public_dtype": str(out.dtype),
            "candidate_dtype": str(candidate_out.dtype),
            "public_device": str(out.device),
            "candidate_device": str(candidate_out.device),
            "public_aliases_input": out.data_ptr() == x.data_ptr(),
            "candidate_aliases_input": candidate_out.data_ptr() == x.data_ptr(),
            "public_metadata_matches_reference": public_meta == ref_meta,
            "candidate_metadata_matches_reference": candidate_meta == ref_meta,
        })
    # Held-out shapes and values are generated independently from the timing
    # fixtures, preventing shape/fixture special-casing from passing silently.
    generator = random.Random(99173)
    for candidate_name, candidate in candidates.items():
      for index in range(10):
        rows_n = generator.randrange(2, 80)
        cols_n = generator.randrange(2, 80)
        if rows_n == cols_n:
            cols_n += 1
        x = torch.randn((rows_n, cols_n), generator=torch.Generator(device="cpu").manual_seed(50000 + index), device="cpu", dtype=torch.float32).to(device)
        before = x.clone()
        ref = polar_express_frozen(x, steps=index % 6)
        out = zeropower_polar_express(x, steps=index % 6)
        candidate_out = candidate(x, steps=index % 6)
        ref_meta = _metadata(ref, x)
        public_meta = _metadata(out, x)
        candidate_meta = _metadata(candidate_out, x)
        rows.append({
            "candidate": candidate_name,
            "case": f"heldout_{index}",
            "public_exact": _same(ref, out),
            "candidate_exact": _same(ref, candidate_out),
            "input_unchanged": _same(x, before),
            "public_shape": list(out.shape),
            "candidate_shape": list(candidate_out.shape),
            "public_dtype": str(out.dtype),
            "candidate_dtype": str(candidate_out.dtype),
            "public_device": str(out.device),
            "candidate_device": str(candidate_out.device),
            "public_aliases_input": out.data_ptr() == x.data_ptr(),
            "candidate_aliases_input": candidate_out.data_ptr() == x.data_ptr(),
            "public_metadata_matches_reference": public_meta == ref_meta,
            "candidate_metadata_matches_reference": candidate_meta == ref_meta,
        })
    static = _static_audit()
    contracts = {name: _contract_audit(zeropower_polar_express, candidate, device) for name, candidate in candidates.items()}
    public_gradients = _public_gradient_audit(zeropower_polar_express, device)
    passed = (
        static["passed"]
        and public_gradients["passed"]
        and all(
            r["public_exact"]
            and r["candidate_exact"]
            and r["input_unchanged"]
            and r["public_metadata_matches_reference"]
            and r["candidate_metadata_matches_reference"]
            for r in rows
        )
        and all(
            check["invalid_steps_rejected"] and check["invalid_rank_rejected"]
            for contract in contracts.values()
            for check in contract.values()
        )
    )
    receipt = {"protocol": "muon-safety-audit-v3", "device": str(device), "passed": passed, "static": static, "public_gradients": public_gradients, "contracts": contracts, "cases": rows, "torch": torch.__version__}
    print(json.dumps(receipt, indent=2))
    if not passed:
        raise SystemExit(1)


if __name__ == "__main__":
    main()

baseline_muon_frozen.py

"""Immutable reference for Muon polar-express evaluation.

This is an independent copy of the pre-candidate kernel.  Evaluators import
this module rather than the mutable production implementation, so changing the
candidate cannot silently change the reference used to judge it.
"""

from __future__ import annotations

import torch


COEFFS = [
    (8.156554524902461, -22.48329292557795, 15.878769915207462),
    (4.042929935166739, -2.808917465908714, 0.5000178451051316),
    (3.8916678022926607, -2.772484153217685, 0.5060648178503393),
    (3.285753657755655, -2.3681294933425376, 0.46449024233003106),
    (2.3465413258596377, -1.7097828382687081, 0.42323551169305323),
]


def polar_express_frozen(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Frozen eager reference, preserving the original operation sequence."""
    assert G.ndim >= 2
    assert steps <= len(COEFFS)

    X = G.bfloat16()
    transpose_needed = G.size(-2) > G.size(-1)
    if transpose_needed:
        X = X.mT

    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-7)

    for a, b, c in COEFFS[:steps]:
        A = X @ X.mT
        A2 = A @ A
        B = b * A + c * A2
        X = a * X + B @ X

    if transpose_needed:
        X = X.mT
    return X

eval_muon_exact.py

"""Strict exact-output evaluator for Muon polar-express candidates.

Promotion requires raw-byte equality for every output in the correctness
matrix, including both matrix orientations, steps 0..5, dtypes, repeated
seeds, and non-contiguous views. Candidates are never timed unless they pass
the exact, contract, and gradient gates.
"""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import platform
import random
import statistics
import sys
import time
from pathlib import Path

import torch


ROOT = Path(__file__).resolve().parents[1]
CANDIDATE_PATH = ROOT / "optimization_candidates" / "muon_exact_candidates.py"
VARIANT_PATH = ROOT / "optimization_candidates" / "muon_compile_variants.py"
FROZEN_PATH = ROOT / "optimization_candidates" / "baseline_muon_frozen.py"


def _load_candidates():
    spec = importlib.util.spec_from_file_location("muon_exact_candidates", CANDIDATE_PATH)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Cannot load candidates at {CANDIDATE_PATH}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    candidates = dict(module.CANDIDATES)
    variant_spec = importlib.util.spec_from_file_location("muon_compile_variants", VARIANT_PATH)
    if variant_spec is None or variant_spec.loader is None:
        raise RuntimeError(f"Cannot load variants at {VARIANT_PATH}")
    variant_module = importlib.util.module_from_spec(variant_spec)
    variant_spec.loader.exec_module(variant_module)
    candidates.update(variant_module.CANDIDATES)
    return candidates


def _load_frozen_baseline():
    spec = importlib.util.spec_from_file_location("baseline_muon_frozen", FROZEN_PATH)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Cannot load frozen baseline at {FROZEN_PATH}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module.polar_express_frozen


def _sync(device):
    if device.type == "mps":
        torch.mps.synchronize()
    elif device.type == "cuda":
        torch.cuda.synchronize(device)


def _make_input(rows, cols, dtype, device, seed, noncontiguous=False):
    gen = torch.Generator(device="cpu").manual_seed(seed)
    base = torch.randn((rows, cols), generator=gen, dtype=torch.float32).to(device=device, dtype=dtype)
    if noncontiguous:
        storage = torch.empty((rows, cols * 2), device=device, dtype=dtype)
        storage[..., ::2] = base
        return storage[..., ::2]
    return base


def _make_batched_input(batch, rows, cols, dtype, device, seed):
    gen = torch.Generator(device="cpu").manual_seed(seed)
    return torch.randn((batch, rows, cols), generator=gen, dtype=torch.float32).to(device=device, dtype=dtype)


def _metadata_signature(tensor: torch.Tensor, input_tensor: torch.Tensor) -> tuple:
    return (
        tuple(tensor.shape),
        tensor.dtype,
        tensor.device.type,
        tensor.device.index,
        tuple(tensor.stride()),
        tensor.storage_offset(),
        tensor.requires_grad,
        tensor.data_ptr() == input_tensor.data_ptr(),
    )


def _bitwise_equal(left: torch.Tensor, right: torch.Tensor) -> bool:
    """Compare exact value bits, including NaN payloads, after metadata check."""
    if left.dtype != right.dtype or left.device.type != right.device.type or left.shape != right.shape:
        return False
    try:
        # ``Tensor.contiguous()`` on MPS can retain a singleton view's
        # unusual stride.  Clone on CPU with an explicit contiguous format
        # before the dtype reinterpretation so the comparator itself cannot
        # reject a valid non-contiguous test input.
        left_bytes = left.detach().cpu().clone(memory_format=torch.contiguous_format).view(torch.uint8)
        right_bytes = right.detach().cpu().clone(memory_format=torch.contiguous_format).view(torch.uint8)
        return torch.equal(left_bytes, right_bytes)
    except (RuntimeError, ValueError):
        return False


def _exact_gate(baseline, candidate, device):
    shapes = [(1, 1), (7, 13), (13, 7), (32, 64), (64, 32), (128, 512), (512, 128)]
    dtypes = [torch.float32, torch.float16, torch.bfloat16]
    rows = []
    for shape_index, (rows_n, cols_n) in enumerate(shapes):
        for dtype_index, dtype in enumerate(dtypes):
            for steps in range(6):
                for noncontiguous in (False, True):
                    seed = 17000 + shape_index * 1000 + dtype_index * 100 + steps * 10 + int(noncontiguous)
                    x = _make_input(rows_n, cols_n, dtype, device, seed, noncontiguous)
                    before = x.clone()
                    reference = baseline(x, steps=steps)
                    output = candidate(x, steps=steps)
                    if not _bitwise_equal(x, before):
                        return {
                            "passed": False,
                            "first_failure": {
                                "shape": [rows_n, cols_n],
                                "dtype": str(dtype),
                                "steps": steps,
                                "noncontiguous": noncontiguous,
                                "reason": "input mutated",
                            },
                            "cases_checked": len(rows),
                        }
                    if _metadata_signature(reference, x) != _metadata_signature(output, x):
                        return {
                            "passed": False,
                            "first_failure": {
                                "shape": [rows_n, cols_n],
                                "dtype": str(dtype),
                                "steps": steps,
                                "noncontiguous": noncontiguous,
                                "reason": "output metadata mismatch",
                                "reference_metadata": repr(_metadata_signature(reference, x)),
                                "candidate_metadata": repr(_metadata_signature(output, x)),
                            },
                            "cases_checked": len(rows),
                        }
                    if not _bitwise_equal(reference, output):
                        return {
                            "passed": False,
                            "first_failure": {
                                "shape": [rows_n, cols_n],
                                "dtype": str(dtype),
                                "steps": steps,
                                "noncontiguous": noncontiguous,
                                "max_abs": (reference - output).abs().max().item(),
                                "differing_elements": int(torch.count_nonzero(reference != output).item()),
                            },
                            "cases_checked": len(rows),
                        }
                    rows.append({"shape": [rows_n, cols_n], "dtype": str(dtype), "steps": steps, "noncontiguous": noncontiguous})
    for dtype_index, dtype in enumerate(dtypes):
        for steps in range(6):
            x = _make_batched_input(2, 7, 13, dtype, device, 22000 + dtype_index * 100 + steps)
            before = x.clone()
            reference = baseline(x, steps=steps)
            output = candidate(x, steps=steps)
            if not _bitwise_equal(x, before) or _metadata_signature(reference, x) != _metadata_signature(output, x) or not _bitwise_equal(reference, output):
                return {
                    "passed": False,
                    "first_failure": {
                        "shape": [2, 7, 13],
                        "dtype": str(dtype),
                        "steps": steps,
                        "noncontiguous": False,
                        "reason": "rank-3 mismatch or input mutation",
                        "max_abs": (reference - output).abs().max().item(),
                    },
                    "cases_checked": len(rows),
                }
            rows.append({"shape": [2, 7, 13], "dtype": str(dtype), "steps": steps, "noncontiguous": False})
    return {"passed": True, "cases_checked": len(rows)}


def _contract_gate(baseline, candidate, device):
    x = _make_input(32, 64, torch.float32, device, 18001)
    # Match the baseline's accepted steps=0 and rejected steps>5 behavior.
    if not _bitwise_equal(baseline(x, steps=0), candidate(x, steps=0)):
        raise AssertionError("steps=0 contract mismatch")
    try:
        candidate(x, steps=6)
    except (AssertionError, RuntimeError):
        pass
    else:
        raise AssertionError("candidate accepted invalid steps=6")
    return {"passed": True, "steps_zero": True, "invalid_steps_rejected": True}


def _gradient_gate(baseline, candidate, device):
    cases = [((13, 7), 5), ((7, 13), 5), ((2, 7, 13), 3), ((32, 64), 1)]
    checked = []
    for shape, steps in cases:
        x0 = torch.randn(shape, device=device, dtype=torch.float32)
        weight = torch.randn(shape, device=device, dtype=torch.bfloat16)
        x_ref = x0.clone().requires_grad_(True)
        x_cand = x0.clone().requires_grad_(True)
        try:
            ref = (baseline(x_ref, steps=steps) * weight).sum()
            out = (candidate(x_cand, steps=steps) * weight).sum()
            ref.backward()
            out.backward()
        except (RuntimeError, AssertionError) as exc:
            # A candidate that only works under no_grad is not silently
            # benchmarked as if it were a general replacement.
            return {
                "passed": False,
                "reason": f"autograd execution failed: {type(exc).__name__}: {exc}",
                "shape": list(shape),
                "steps": steps,
            }
        if not _bitwise_equal(x_ref.grad, x_cand.grad):
            return {
                "passed": False,
                "reason": "gradient bits differ",
                "shape": list(shape),
                "steps": steps,
                "max_abs": (x_ref.grad - x_cand.grad).abs().max().item(),
            }
        checked.append({"shape": list(shape), "steps": steps, "grad_dtype": str(x_cand.grad.dtype)})
    return {"passed": True, "cases_checked": checked}


def _timed(fn, x, steps):
    _sync(x.device)
    start = time.perf_counter_ns()
    fn(x, steps=steps)
    _sync(x.device)
    return (time.perf_counter_ns() - start) / 1e6


def _benchmark(baseline, candidate, device, warmup, repeats):
    specs = [
        (128, 512, torch.float32, 5, 19001),
        (512, 128, torch.float32, 5, 19002),
        (256, 1024, torch.float32, 3, 19003),
        # Shapes of the kit's d_model/d_ff projection parameters.
        (512, 2048, torch.float32, 5, 19004),
        (2048, 512, torch.float32, 5, 19005),
        (1536, 512, torch.float32, 5, 19006),
    ]
    rows = []
    for index, (rows_n, cols_n, dtype, steps, seed) in enumerate(specs):
        x = _make_input(rows_n, cols_n, dtype, device, seed)
        before_timing = x.detach().clone(memory_format=torch.contiguous_format)
        for _ in range(warmup):
            baseline(x, steps=steps)
            candidate(x, steps=steps)
        _sync(device)
        baseline_ms, candidate_ms = [], []
        order = random.Random(19500 + index)
        for _ in range(repeats):
            if order.randrange(2):
                candidate_ms.append(_timed(candidate, x, steps))
                baseline_ms.append(_timed(baseline, x, steps))
            else:
                baseline_ms.append(_timed(baseline, x, steps))
                candidate_ms.append(_timed(candidate, x, steps))
        baseline_median = statistics.median(baseline_ms)
        candidate_median = statistics.median(candidate_ms)
        input_unchanged = _bitwise_equal(x, before_timing)
        if not input_unchanged:
            raise AssertionError(f"candidate mutated benchmark input for shape={(rows_n, cols_n)}")
        rows.append({
            "shape": [rows_n, cols_n],
            "dtype": str(dtype),
            "steps": steps,
            "baseline_ms": baseline_ms,
            "candidate_ms": candidate_ms,
            "baseline_median_ms": baseline_median,
            "candidate_median_ms": candidate_median,
            "speedup": baseline_median / candidate_median,
            "baseline_p90_ms": statistics.quantiles(baseline_ms, n=10)[8],
            "candidate_p90_ms": statistics.quantiles(candidate_ms, n=10)[8],
            "input_unchanged": input_unchanged,
        })
    return rows


def _hash(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", choices=("auto", "cpu", "mps"), default="auto")
    parser.add_argument("--candidate", choices=("inplace", "addmm", "inplace_add", "preallocated_mm", "contiguous_transpose_mm", "contiguous_tall_only", "cast_after_transpose_tall", "clone_final_tall", "out_mm_bx_only", "out_mm_a2_only", "out_mm_bx_muon_context", "out_mm_bx_cast_after_context", "aot_eager", "inductor_no_fusion", "elementwise_only", "eager_mm", "all"), default="all")
    parser.add_argument("--warmup", type=int, default=12)
    parser.add_argument("--repeats", type=int, default=40)
    parser.add_argument("--out", type=Path)
    args = parser.parse_args()
    if args.device == "mps" and not torch.backends.mps.is_available():
        raise SystemExit("MPS requested but unavailable")
    device = torch.device("mps" if args.device == "auto" and torch.backends.mps.is_available() else ("cpu" if args.device == "auto" else args.device))
    all_candidates = _load_candidates()
    selected = all_candidates if args.candidate == "all" else {args.candidate: all_candidates[args.candidate]}
    receipt = {
        "protocol": "muon-polar-exact-v1",
        "exact_rule": "raw tensor bytes for every correctness case; no tolerance promotion",
        "device": str(device),
        "platform": platform.platform(),
        "python": sys.version,
        "torch": torch.__version__,
        "source_hashes": {
            "baseline_muon": _hash(FROZEN_PATH),
            "production_muon": _hash(ROOT / "optimizers" / "muon.py"),
            "candidates": _hash(CANDIDATE_PATH),
            "compile_variants": _hash(VARIANT_PATH),
            "evaluator": _hash(Path(__file__)),
        },
        "candidates": {},
    }
    print(f"device={device} torch={torch.__version__}")
    frozen_baseline = _load_frozen_baseline()
    for name, candidate in selected.items():
        print(f"checking={name}")
        exact = _exact_gate(frozen_baseline, candidate, device)
        result = {"exact_gate": exact}
        if exact["passed"]:
            result["contract_gate"] = _contract_gate(frozen_baseline, candidate, device)
            result["gradient_gate"] = _gradient_gate(frozen_baseline, candidate, device)
            if result["contract_gate"].get("passed") and result["gradient_gate"].get("passed"):
                result["benchmark"] = _benchmark(frozen_baseline, candidate, device, args.warmup, args.repeats)
                for row in result["benchmark"]:
                    print(f"{name} shape={row['shape']} baseline={row['baseline_median_ms']:.4f}ms candidate={row['candidate_median_ms']:.4f}ms speedup={row['speedup']:.3f}x")
            else:
                result["benchmark"] = {"skipped": True, "reason": "contract or gradient gate failed"}
        else:
            print(f"{name} REJECTED exact={exact['first_failure']}")
        receipt["candidates"][name] = result
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(receipt, indent=2) + "\n")
        print(f"receipt={args.out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

eval_muon_public_api.py

"""Public-API confirmation for the MPS Muon polar-express optimization."""

from __future__ import annotations

import hashlib
import json
import random
import statistics
import sys
import time
from pathlib import Path

import torch

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "optimization_candidates"))
from baseline_muon_frozen import polar_express_frozen

from optimizers.muon import zeropower_polar_express


def _sync(device):
    if device.type == "mps":
        torch.mps.synchronize()
    elif device.type == "cuda":
        torch.cuda.synchronize(device)


def _make(shape, seed, device):
    generator = torch.Generator(device="cpu").manual_seed(seed)
    return torch.randn(shape, generator=generator, dtype=torch.float32).to(device)


def _metadata(tensor, source):
    return (
        tuple(tensor.shape), tensor.dtype, tensor.device.type, tensor.device.index,
        tuple(tensor.stride()), tensor.storage_offset(), tensor.requires_grad,
        tensor.data_ptr() == source.data_ptr(),
    )


def _bitwise_equal(a, b):
    if a.shape != b.shape or a.dtype != b.dtype or a.device.type != b.device.type:
        return False
    try:
        return torch.equal(
            a.detach().cpu().clone(memory_format=torch.contiguous_format).view(torch.uint8),
            b.detach().cpu().clone(memory_format=torch.contiguous_format).view(torch.uint8),
        )
    except (RuntimeError, ValueError):
        return False


def _timed(fn, tensors, device):
    _sync(device)
    start = time.perf_counter_ns()
    for tensor in tensors:
        with torch.no_grad():
            fn(tensor, steps=5)
    _sync(device)
    return (time.perf_counter_ns() - start) / 1e6


def main():
    device = torch.device("mps")
    if not torch.backends.mps.is_available():
        raise SystemExit("MPS unavailable")
    shapes = [(1536, 512), (2048, 512), (512, 2048), (512, 512)]
    baseline_inputs = [_make(shape, 61000 + i, device) for i, shape in enumerate(shapes)]
    public_inputs = [tensor.clone() for tensor in baseline_inputs]
    exact_rows = []
    for index, (base_input, public_input) in enumerate(zip(baseline_inputs, public_inputs)):
        base_before = base_input.clone()
        public_before = public_input.clone()
        reference = polar_express_frozen(base_input, steps=5)
        output = zeropower_polar_express(public_input, steps=5)
        exact_rows.append({
            "shape": list(base_input.shape),
            "values_exact": _bitwise_equal(reference, output),
            "metadata_exact": _metadata(reference, base_input) == _metadata(output, public_input),
            "baseline_input_unchanged": _bitwise_equal(base_input, base_before),
            "public_input_unchanged": _bitwise_equal(public_input, public_before),
        })
        if not all(exact_rows[-1][key] for key in ("values_exact", "metadata_exact", "baseline_input_unchanged", "public_input_unchanged")):
            raise AssertionError(f"public API mismatch for {tuple(base_input.shape)}")
    for _ in range(24):
        _timed(polar_express_frozen, baseline_inputs, device)
        _timed(zeropower_polar_express, public_inputs, device)
    baseline_ms, public_ms = [], []
    order = random.Random(62001)
    for _ in range(240):
        if order.randrange(2):
            public_ms.append(_timed(zeropower_polar_express, public_inputs, device))
            baseline_ms.append(_timed(polar_express_frozen, baseline_inputs, device))
        else:
            baseline_ms.append(_timed(polar_express_frozen, baseline_inputs, device))
            public_ms.append(_timed(zeropower_polar_express, public_inputs, device))
    baseline_median = statistics.median(baseline_ms)
    public_median = statistics.median(public_ms)
    receipt = {
        "protocol": "muon-public-api-v1",
        "device": str(device),
        "torch": torch.__version__,
        "shapes": [list(shape) for shape in shapes],
        "exact_rows": exact_rows,
        "baseline_sha256": hashlib.sha256((ROOT / "optimization_candidates" / "baseline_muon_frozen.py").read_bytes()).hexdigest(),
        "production_sha256": hashlib.sha256((ROOT / "optimizers" / "muon.py").read_bytes()).hexdigest(),
        "baseline_median_ms": baseline_median,
        "public_median_ms": public_median,
        "speedup": baseline_median / public_median,
        "baseline_p90_ms": statistics.quantiles(baseline_ms, n=10)[8],
        "public_p90_ms": statistics.quantiles(public_ms, n=10)[8],
        "baseline_ms": baseline_ms,
        "public_ms": public_ms,
    }
    out = ROOT / "optimization_candidates" / "results" / "muon-public-api-final.json"
    out.write_text(json.dumps(receipt, indent=2) + "\n")
    print(f"device={device} exact=True baseline={baseline_median:.4f}ms public={public_median:.4f}ms speedup={baseline_median / public_median:.3f}x")
    print(f"receipt={out}")


if __name__ == "__main__":
    main()

muon-public-api-final.json

{
  "protocol": "muon-public-api-v1",
  "device": "mps",
  "torch": "2.11.0",
  "shapes": [
    [
      1536,
      512
    ],
    [
      2048,
      512
    ],
    [
      512,
      2048
    ],
    [
      512,
      512
    ]
  ],
  "exact_rows": [
    {
      "shape": [
        1536,
        512
      ],
      "values_exact": true,
      "metadata_exact": true,
      "baseline_input_unchanged": true,
      "public_input_unchanged": true
    },
    {
      "shape": [
        2048,
        512
      ],
      "values_exact": true,
      "metadata_exact": true,
      "baseline_input_unchanged": true,
      "public_input_unchanged": true
    },
    {
      "shape": [
        512,
        2048
      ],
      "values_exact": true,
      "metadata_exact": true,
      "baseline_input_unchanged": true,
      "public_input_unchanged": true
    },
    {
      "shape": [
        512,
        512
      ],
      "values_exact": true,
      "metadata_exact": true,
      "baseline_input_unchanged": true,
      "public_input_unchanged": true
    }
  ],
  "baseline_sha256": "0f9fa3067f272781f8652b3b60eabd513c3750c9e299cd23fc1f6a8934ba6e14",
  "production_sha256": "c82e33364908724ea23d02676189612a98b3e516ef4b28732e633d6914e088f5",
  "baseline_median_ms": 17.490353499999998,
  "public_median_ms": 16.0311455,
  "speedup": 1.0910233145847248,
  "baseline_p90_ms": 18.610879399999998,
  "public_p90_ms": 17.4119706,
  "baseline_ms": [
    19.169958,
    17.1085,
    16.643791,
    16.406916,
    17.88825,
    17.305917,
    16.761208,
    16.364,
    16.637625,
    18.158125,
    18.299834,
    18.782458,
    17.814875,
    17.03875,
    17.454666,
    17.900708,
    16.90575,
    17.8685,
    16.64725,
    17.235208,
    17.500333,
    18.789209,
    16.74175,
    18.010167,
    17.214625,
    17.931583,
    16.899625,
    16.883625,
    16.065083,
    16.359,
    16.871875,
    17.619292,
    17.178,
    17.44175,
    17.229708,
    17.318792,
    16.469541,
    16.370459,
    17.167667,
    17.607334,
    17.080792,
    16.881084,
    18.132958,
    18.096542,
    17.14525,
    17.099291,
    17.029375,
    18.323541,
    16.278709,
    17.18425,
    18.189,
    17.060792,
    16.99225,
    17.656833,
    16.907375,
    19.867083,
    17.035625,
    16.421958,
    17.211167,
    17.600625,
    16.694542,
    17.986208,
    17.553709,
    17.047167,
    17.600625,
    18.918833,
    16.332208,
    19.151166,
    17.097208,
    17.726709,
    17.482916,
    18.083417,
    17.531208,
    17.520041,
    17.517666,
    16.049417,
    16.077833,
    17.389708,
    17.12775,
    17.288959,
    16.741917,
    16.672709,
    16.226667,
    18.235792,
    16.586375,
    16.897625,
    16.737625,
    16.733917,
    18.20525,
    16.463667,
    17.741875,
    17.276084,
    17.515166,
    18.062334,
    17.065084,
    16.836875,
    17.265542,
    17.644917,
    16.898666,
    17.204709,
    19.241458,
    16.34725,
    16.439167,
    17.497791,
    17.318,
    17.172625,
    17.184208,
    18.171167,
    17.097459,
    17.720833,
    17.011792,
    18.066084,
    17.209167,
    18.1835,
    16.900709,
    17.268708,
    16.9655,
    17.815334,
    16.86975,
    17.530667,
    17.6785,
    18.486625,
    16.87975,
    17.857292,
    17.181792,
    16.819583,
    17.943666,
    17.284459,
    17.663875,
    17.102125,
    17.466917,
    17.429167,
    17.566875,
    17.157792,
    17.371167,
    17.433833,
    16.840709,
    16.664458,
    17.701583,
    16.883375,
    16.710209,
    16.925834,
    18.893,
    17.809416,
    18.785083,
    18.464625,
    16.924916,
    18.168667,
    17.942083,
    16.8305,
    16.803917,
    17.957334,
    19.301875,
    17.89375,
    17.649875,
    18.013833,
    17.617834,
    16.80175,
    17.595667,
    17.422083,
    16.731458,
    17.829375,
    17.31525,
    17.203459,
    18.604916,
    18.611542,
    17.587042,
    18.071292,
    17.906625,
    17.683,
    17.606916,
    19.305959,
    17.454625,
    17.789584,
    17.912625,
    17.312666,
    18.509583,
    17.641625,
    16.940084,
    18.210625,
    18.57225,
    16.776708,
    18.18,
    18.89975,
    17.721,
    17.234667,
    18.630625,
    18.014292,
    18.197542,
    18.871667,
    17.922792,
    17.743875,
    17.521416,
    16.894375,
    18.530208,
    17.578791,
    24.268917,
    17.249333,
    20.345125,
    19.867875,
    17.428833,
    18.424417,
    17.330459,
    17.375208,
    17.77825,
    16.460709,
    18.486,
    18.568417,
    19.482625,
    18.723292,
    19.350333,
    17.954417,
    17.655917,
    17.835584,
    19.610833,
    17.919666,
    17.480042,
    19.975959,
    17.853667,
    16.735333,
    17.250458,
    17.251167,
    16.958333,
    18.102167,
    17.078334,
    17.757,
    18.265125,
    18.231792,
    17.2205,
    18.217167,
    17.561042,
    17.414542,
    18.415167,
    16.910708,
    17.968167,
    18.616417,
    18.185375,
    17.339375,
    17.263958,
    17.761083
  ],
  "public_ms": [
    17.788625,
    16.437542,
    16.3735,
    15.919917,
    16.556,
    16.304084,
    16.159958,
    15.20275,
    16.123541,
    15.44375,
    16.838542,
    15.686916,
    16.718333,
    15.596292,
    15.772,
    15.520125,
    15.997541,
    15.84,
    16.307042,
    15.149375,
    17.591209,
    17.060709,
    17.23,
    15.288583,
    17.023542,
    16.422917,
    16.32775,
    17.103333,
    16.998,
    16.012167,
    15.018375,
    14.934791,
    15.044959,
    15.230125,
    15.058,
    15.44925,
    15.802208,
    15.229,
    16.928333,
    16.393583,
    15.638833,
    16.60975,
    16.27825,
    16.031625,
    16.157458,
    15.3625,
    15.525583,
    15.890125,
    15.743334,
    16.191417,
    16.169458,
    16.215042,
    15.3255,
    16.249667,
    15.666042,
    16.19325,
    15.0425,
    15.312042,
    15.88,
    17.065792,
    16.395166,
    15.330167,
    15.3365,
    15.408917,
    16.623917,
    15.698125,
    15.36825,
    16.405917,
    17.213667,
    18.49175,
    15.291833,
    15.801542,
    16.071375,
    17.555875,
    15.126,
    16.921042,
    15.897416,
    15.544541,
    15.000667,
    14.863333,
    15.675791,
    15.714208,
    15.481666,
    15.699583,
    16.492541,
    15.492791,
    16.212041,
    15.702583,
    16.853625,
    15.840917,
    16.196875,
    15.475959,
    16.393375,
    16.327416,
    15.949666,
    15.9575,
    16.069625,
    15.931208,
    15.843,
    15.705584,
    16.026167,
    16.538917,
    16.761209,
    14.835625,
    16.045292,
    15.863959,
    15.801125,
    15.342416,
    16.191958,
    15.635708,
    16.392458,
    15.597916,
    15.560625,
    16.166875,
    15.423792,
    16.35975,
    16.233,
    15.548625,
    15.454459,
    15.510333,
    15.496459,
    16.500083,
    15.118125,
    15.936,
    16.848917,
    16.441333,
    17.338333,
    16.648083,
    17.404209,
    17.412833,
    16.111667,
    16.611958,
    17.472375,
    14.867459,
    16.006,
    15.694792,
    16.399209,
    16.83325,
    15.888208,
    17.463666,
    17.674958,
    17.68275,
    15.43575,
    18.133958,
    15.940292,
    15.169917,
    16.11225,
    14.99275,
    15.34825,
    16.066208,
    16.416375,
    14.943167,
    17.145458,
    16.603666,
    15.874958,
    15.385083,
    15.367167,
    17.551875,
    15.629834,
    15.552208,
    17.123125,
    15.886417,
    16.210792,
    16.030666,
    15.115042,
    15.842291,
    15.88175,
    15.580417,
    16.690041,
    17.480042,
    18.207875,
    16.374667,
    16.695041,
    16.914625,
    15.9235,
    16.454459,
    16.436792,
    15.955541,
    17.006209,
    15.558792,
    16.807625,
    18.536459,
    15.663,
    16.062708,
    17.653583,
    16.923959,
    16.509208,
    16.382416,
    16.168166,
    16.383958,
    15.286917,
    16.3795,
    15.4015,
    19.44075,
    20.918292,
    20.6095,
    20.298375,
    16.939542,
    16.656375,
    16.581167,
    16.863334,
    16.173125,
    16.08925,
    15.011791,
    15.05975,
    16.563333,
    15.521209,
    18.012167,
    16.407291,
    16.995209,
    17.845666,
    18.390458,
    17.399375,
    18.184625,
    16.355666,
    17.473375,
    16.721958,
    15.639459,
    16.314625,
    15.813125,
    15.231584,
    15.795958,
    15.515625,
    15.025334,
    15.672459,
    15.716084,
    15.433291,
    15.070541,
    15.800042,
    15.763334,
    15.850541,
    15.417834,
    15.288375,
    16.165334,
    15.312083,
    14.848,
    15.03175,
    17.045625,
    15.270958,
    15.497125
  ]
}

muon-safety-final-v6.json

{
  "protocol": "muon-safety-audit-v3",
  "device": "mps",
  "passed": true,
  "static": {
    "passed": true,
    "findings": []
  },
  "public_gradients": {
    "passed": true,
    "cases": [
      {
        "shape": [
          13,
          7
        ],
        "steps": 5,
        "gradients_exact": true,
        "execution_error": null
      },
      {
        "shape": [
          7,
          13
        ],
        "steps": 5,
        "gradients_exact": true,
        "execution_error": null
      },
      {
        "shape": [
          2,
          7,
          13
        ],
        "steps": 3,
        "gradients_exact": true,
        "execution_error": null
      },
      {
        "shape": [
          32,
          64
        ],
        "steps": 1,
        "gradients_exact": true,
        "execution_error": null
      }
    ]
  },
  "contracts": {
    "contiguous_tall_only": {
      "public": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      },
      "candidate": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      }
    },
    "cast_after_transpose_tall": {
      "public": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      },
      "candidate": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      }
    },
    "clone_final_tall": {
      "public": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      },
      "candidate": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      }
    },
    "out_mm_bx_only": {
      "public": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      },
      "candidate": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      }
    },
    "out_mm_a2_only": {
      "public": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      },
      "candidate": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      }
    },
    "out_mm_bx_muon_context": {
      "public": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      },
      "candidate": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      }
    },
    "out_mm_bx_cast_after_context": {
      "public": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      },
      "candidate": {
        "invalid_steps_rejected": true,
        "invalid_rank_rejected": true,
        "steps_zero_shape": [
          1536,
          512
        ],
        "negative_steps_shape": [
          1536,
          512
        ]
      }
    }
  },
  "cases": [
    {
      "candidate": "contiguous_tall_only",
      "case": "zeros",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "ones",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "large",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "small",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "mixed",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "pattern",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "noncontiguous",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "zeros",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "ones",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "large",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "small",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "mixed",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "pattern",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "noncontiguous",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "zeros",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "ones",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "large",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "small",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "mixed",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "pattern",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "noncontiguous",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "zeros",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "ones",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "large",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "small",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "mixed",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "pattern",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "noncontiguous",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "zeros",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "ones",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "large",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "small",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "mixed",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "pattern",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "noncontiguous",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "zeros",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "ones",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "large",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "small",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "mixed",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "pattern",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "noncontiguous",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "zeros",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "ones",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "large",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "small",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "mixed",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "pattern",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "noncontiguous",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        1536,
        512
      ],
      "candidate_shape": [
        1536,
        512
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "heldout_0",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        25,
        49
      ],
      "candidate_shape": [
        25,
        49
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "heldout_1",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        38,
        19
      ],
      "candidate_shape": [
        38,
        19
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "heldout_2",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        24,
        70
      ],
      "candidate_shape": [
        24,
        70
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "heldout_3",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        38,
        15
      ],
      "candidate_shape": [
        38,
        15
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "heldout_4",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        70,
        13
      ],
      "candidate_shape": [
        70,
        13
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "heldout_5",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        32,
        21
      ],
      "candidate_shape": [
        32,
        21
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "heldout_6",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        23,
        35
      ],
      "candidate_shape": [
        23,
        35
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "heldout_7",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        68,
        34
      ],
      "candidate_shape": [
        68,
        34
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "heldout_8",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        78,
        16
      ],
      "candidate_shape": [
        78,
        16
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "contiguous_tall_only",
      "case": "heldout_9",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        45,
        79
      ],
      "candidate_shape": [
        45,
        79
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "heldout_0",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        66,
        13
      ],
      "candidate_shape": [
        66,
        13
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "heldout_1",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        11,
        13
      ],
      "candidate_shape": [
        11,
        13
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "heldout_2",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        54,
        47
      ],
      "candidate_shape": [
        54,
        47
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "heldout_3",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        70,
        56
      ],
      "candidate_shape": [
        70,
        56
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "heldout_4",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        37,
        58
      ],
      "candidate_shape": [
        37,
        58
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "heldout_5",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        27,
        43
      ],
      "candidate_shape": [
        27,
        43
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "heldout_6",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        8,
        50
      ],
      "candidate_shape": [
        8,
        50
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "heldout_7",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        68,
        27
      ],
      "candidate_shape": [
        68,
        27
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "heldout_8",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        48,
        23
      ],
      "candidate_shape": [
        48,
        23
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "cast_after_transpose_tall",
      "case": "heldout_9",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        77,
        20
      ],
      "candidate_shape": [
        77,
        20
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "heldout_0",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        68,
        35
      ],
      "candidate_shape": [
        68,
        35
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "heldout_1",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        71,
        66
      ],
      "candidate_shape": [
        71,
        66
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "heldout_2",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        57,
        42
      ],
      "candidate_shape": [
        57,
        42
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "heldout_3",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        48,
        36
      ],
      "candidate_shape": [
        48,
        36
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "heldout_4",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        38,
        12
      ],
      "candidate_shape": [
        38,
        12
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "heldout_5",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        7,
        79
      ],
      "candidate_shape": [
        7,
        79
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "heldout_6",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        7,
        8
      ],
      "candidate_shape": [
        7,
        8
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "heldout_7",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        11,
        16
      ],
      "candidate_shape": [
        11,
        16
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "heldout_8",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        78,
        51
      ],
      "candidate_shape": [
        78,
        51
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "clone_final_tall",
      "case": "heldout_9",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        34,
        63
      ],
      "candidate_shape": [
        34,
        63
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "heldout_0",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        13,
        2
      ],
      "candidate_shape": [
        13,
        2
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "heldout_1",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        21,
        22
      ],
      "candidate_shape": [
        21,
        22
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "heldout_2",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        35,
        17
      ],
      "candidate_shape": [
        35,
        17
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "heldout_3",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        46,
        27
      ],
      "candidate_shape": [
        46,
        27
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "heldout_4",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        39,
        28
      ],
      "candidate_shape": [
        39,
        28
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "heldout_5",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        53,
        77
      ],
      "candidate_shape": [
        53,
        77
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "heldout_6",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        76,
        64
      ],
      "candidate_shape": [
        76,
        64
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "heldout_7",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        38,
        40
      ],
      "candidate_shape": [
        38,
        40
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "heldout_8",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        44,
        78
      ],
      "candidate_shape": [
        44,
        78
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_only",
      "case": "heldout_9",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        5,
        12
      ],
      "candidate_shape": [
        5,
        12
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "heldout_0",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        16,
        73
      ],
      "candidate_shape": [
        16,
        73
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "heldout_1",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        35,
        13
      ],
      "candidate_shape": [
        35,
        13
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "heldout_2",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        72,
        9
      ],
      "candidate_shape": [
        72,
        9
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "heldout_3",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        11,
        21
      ],
      "candidate_shape": [
        11,
        21
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "heldout_4",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        71,
        2
      ],
      "candidate_shape": [
        71,
        2
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "heldout_5",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        23,
        40
      ],
      "candidate_shape": [
        23,
        40
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "heldout_6",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        29,
        75
      ],
      "candidate_shape": [
        29,
        75
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "heldout_7",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        47,
        73
      ],
      "candidate_shape": [
        47,
        73
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "heldout_8",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        19,
        49
      ],
      "candidate_shape": [
        19,
        49
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_a2_only",
      "case": "heldout_9",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        49,
        2
      ],
      "candidate_shape": [
        49,
        2
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "heldout_0",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        52,
        14
      ],
      "candidate_shape": [
        52,
        14
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "heldout_1",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        67,
        6
      ],
      "candidate_shape": [
        67,
        6
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "heldout_2",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        7,
        16
      ],
      "candidate_shape": [
        7,
        16
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "heldout_3",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        78,
        11
      ],
      "candidate_shape": [
        78,
        11
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "heldout_4",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        19,
        26
      ],
      "candidate_shape": [
        19,
        26
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "heldout_5",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        17,
        27
      ],
      "candidate_shape": [
        17,
        27
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "heldout_6",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        47,
        49
      ],
      "candidate_shape": [
        47,
        49
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "heldout_7",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        70,
        51
      ],
      "candidate_shape": [
        70,
        51
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "heldout_8",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        57,
        53
      ],
      "candidate_shape": [
        57,
        53
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_muon_context",
      "case": "heldout_9",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        68,
        77
      ],
      "candidate_shape": [
        68,
        77
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "heldout_0",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        58,
        25
      ],
      "candidate_shape": [
        58,
        25
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "heldout_1",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        71,
        77
      ],
      "candidate_shape": [
        71,
        77
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "heldout_2",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        13,
        43
      ],
      "candidate_shape": [
        13,
        43
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "heldout_3",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        42,
        25
      ],
      "candidate_shape": [
        42,
        25
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "heldout_4",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        27,
        71
      ],
      "candidate_shape": [
        27,
        71
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "heldout_5",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        29,
        66
      ],
      "candidate_shape": [
        29,
        66
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "heldout_6",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        43,
        21
      ],
      "candidate_shape": [
        43,
        21
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "heldout_7",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        17,
        34
      ],
      "candidate_shape": [
        17,
        34
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "heldout_8",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        3,
        10
      ],
      "candidate_shape": [
        3,
        10
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    },
    {
      "candidate": "out_mm_bx_cast_after_context",
      "case": "heldout_9",
      "public_exact": true,
      "candidate_exact": true,
      "input_unchanged": true,
      "public_shape": [
        9,
        56
      ],
      "candidate_shape": [
        9,
        56
      ],
      "public_dtype": "torch.bfloat16",
      "candidate_dtype": "torch.bfloat16",
      "public_device": "mps:0",
      "candidate_device": "mps:0",
      "public_aliases_input": false,
      "candidate_aliases_input": false,
      "public_metadata_matches_reference": true,
      "candidate_metadata_matches_reference": true
    }
  ],
  "torch": "2.11.0"
}

muon_exact_candidates.py

"""Exact-output Muon polar-express kernel candidates.

Every candidate is intentionally independent from the frozen implementation.
The evaluator selects candidates by name and rejects any non-bitwise-equal
output before timing it.
"""

from __future__ import annotations

import torch


COEFFS = [
    (8.156554524902461, -22.48329292557795, 15.878769915207462),
    (4.042929935166739, -2.808917465908714, 0.5000178451051316),
    (3.8916678022926607, -2.772484153217685, 0.5060648178503393),
    (3.285753657755655, -2.3681294933425376, 0.46449024233003106),
    (2.3465413258596377, -1.7097828382687081, 0.42323551169305323),
]


def _prepare(G: torch.Tensor, steps: int):
    assert G.ndim >= 2
    assert steps <= len(COEFFS)
    X = G.bfloat16()
    transpose_needed = G.size(-2) > G.size(-1)
    if transpose_needed:
        X = X.mT
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-7)
    return X, transpose_needed


def baseline_order_inplace(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Reuse B's storage for the polynomial combination, same op sequence."""
    X, transpose_needed = _prepare(G, steps)
    for a, b, c in COEFFS[:steps]:
        A = X @ X.mT
        A2 = A @ A
        # A is dead after this point; reusing its storage removes B allocation.
        A.mul_(b).add_(A2, alpha=c)
        scaled_X = X * a
        X = scaled_X + (A @ X)
    if transpose_needed:
        X = X.mT.contiguous()
    return X


def baseline_order_addmm(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Use addmm for the final affine matrix operation, no compile/fusion."""
    X, transpose_needed = _prepare(G, steps)
    for a, b, c in COEFFS[:steps]:
        A = X @ X.mT
        A2 = A @ A
        B = b * A + c * A2
        X = torch.addmm(X * a, B, X) if X.ndim == 2 else (X * a + B @ X)
    if transpose_needed:
        X = X.mT
    return X


def baseline_order_inplace_add(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Reuse B storage and perform the final add in-place on scaled X."""
    X, transpose_needed = _prepare(G, steps)
    for a, b, c in COEFFS[:steps]:
        A = X @ X.mT
        A2 = A @ A
        A.mul_(b).add_(A2, alpha=c)
        scaled_X = X * a
        scaled_X.add_(A @ X)
        X = scaled_X
    if transpose_needed:
        X = X.mT
    return X


def preallocated_mm(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """2-D Muon path using persistent ``out=`` workspaces.

    The mathematical and dispatch order is kept identical to the baseline:
    ``mm, mm, mul/add, mm, mul/add``.  Only temporary storage is reused across
    iterations, which targets MPS allocation/command-encoding overhead.
    """
    X, transpose_needed = _prepare(G, steps)
    if X.ndim != 2 or steps == 0:
        for a, b, c in COEFFS[:steps]:
            A = X @ X.mT
            A2 = A @ A
            B = b * A + c * A2
            X = a * X + B @ X
        if transpose_needed:
            X = X.mT
        return X

    rows = X.shape[-2]
    cols = X.shape[-1]
    A = torch.empty((rows, rows), device=X.device, dtype=X.dtype)
    A2 = torch.empty_like(A)
    B = torch.empty_like(A)
    BX = torch.empty_like(X)
    X_next = torch.empty_like(X)
    for a, b, c in COEFFS[:steps]:
        torch.mm(X, X.mT, out=A)
        torch.mm(A, A, out=A2)
        torch.mul(A, b, out=B)
        torch.add(B, A2, alpha=c, out=B)
        torch.mm(B, X, out=BX)
        torch.mul(X, a, out=X_next)
        torch.add(X_next, BX, out=X_next)
        X, X_next = X_next, X
    if transpose_needed:
        X = X.mT
    return X


def contiguous_transpose_mm(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Exact explicit-mm path with contiguous storage for tall gradients."""
    assert G.ndim >= 2
    assert steps <= len(COEFFS)
    X = G.bfloat16()
    transpose_needed = G.size(-2) > G.size(-1)
    if transpose_needed:
        X = X.mT.contiguous()
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-7)
    for a, b, c in COEFFS[:steps]:
        if X.ndim == 2:
            A = torch.mm(X, X.mT)
            A2 = torch.mm(A, A)
            B = b * A + c * A2
            X = a * X + torch.mm(B, X)
        else:
            A = X @ X.mT
            A2 = A @ A
            B = b * A + c * A2
            X = a * X + B @ X
    if transpose_needed:
        # The frozen baseline's transpose view is followed by a final
        # contiguous result in the public API.  Preserve that metadata contract
        # while keeping the optimized internal layout for the matmuls.
        X = X.mT.contiguous()
    return X


def contiguous_tall_only(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Use the contiguous/explicit-mm path only for tall 2-D matrices.

    Muon's LLM projection weights are mostly either ``d_ff x d_model`` (tall)
    or ``d_model x d_ff`` (wide).  The tall path pays for a strided transpose
    on MPS; the wide path is left in the baseline operation form to avoid
    trading away the win on the opposite orientation.
    """
    assert G.ndim >= 2
    assert steps <= len(COEFFS)
    X = G.bfloat16()
    transpose_needed = G.size(-2) > G.size(-1)
    use_tall_2d_path = X.ndim == 2 and transpose_needed
    if transpose_needed:
        X = X.mT.contiguous() if use_tall_2d_path else X.mT
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-7)
    for a, b, c in COEFFS[:steps]:
        if use_tall_2d_path:
            A = torch.mm(X, X.mT)
            A2 = torch.mm(A, A)
            B = b * A + c * A2
            X = a * X + torch.mm(B, X)
        else:
            A = X @ X.mT
            A2 = A @ A
            B = b * A + c * A2
            X = a * X + B @ X
    if transpose_needed:
        X = X.mT.contiguous()
    return X


def cast_after_transpose_tall(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Cast after materializing the transposed tall input."""
    assert G.ndim >= 2
    assert steps <= len(COEFFS)
    transpose_needed = G.size(-2) > G.size(-1)
    if transpose_needed and G.ndim == 2:
        # Compared with the promoted path, this changes only the order of the
        # layout conversion and dtype cast.  The post-cast tensor is contiguous
        # and all subsequent arithmetic is unchanged.
        X = G.mT.contiguous().bfloat16()
    else:
        X = G.bfloat16()
        if transpose_needed:
            X = X.mT
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-7)
    for a, b, c in COEFFS[:steps]:
        if transpose_needed and G.ndim == 2:
            A = torch.mm(X, X.mT)
            A2 = torch.mm(A, A)
            B = b * A + c * A2
            X = a * X + torch.mm(B, X)
        else:
            A = X @ X.mT
            A2 = A @ A
            B = b * A + c * A2
            X = a * X + B @ X
    if transpose_needed:
        X = X.mT.contiguous() if G.ndim == 2 else X.mT
    return X


def clone_final_tall(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Use clone(memory_format=contiguous) for the required final layout."""
    assert G.ndim >= 2
    assert steps <= len(COEFFS)
    transpose_needed = G.size(-2) > G.size(-1)
    X = G.bfloat16()
    if transpose_needed:
        X = X.mT.contiguous()
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-7)
    for a, b, c in COEFFS[:steps]:
        if transpose_needed and G.ndim == 2:
            A = torch.mm(X, X.mT)
            A2 = torch.mm(A, A)
            B = b * A + c * A2
            X = a * X + torch.mm(B, X)
        else:
            A = X @ X.mT
            A2 = A @ A
            B = b * A + c * A2
            X = a * X + B @ X
    if transpose_needed:
        X = X.mT.clone(memory_format=torch.contiguous_format)
    return X


def out_mm_bx_only(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Reuse only the B@X result buffer; retain baseline scalar operations."""
    assert G.ndim >= 2
    assert steps <= len(COEFFS)
    transpose_needed = G.size(-2) > G.size(-1)
    tall = transpose_needed and G.ndim == 2
    X = G.bfloat16()
    if tall:
        X = X.mT.contiguous()
    elif transpose_needed:
        X = X.mT
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-7)
    bx_buffer = torch.empty_like(X) if tall else None
    for a, b, c in COEFFS[:steps]:
        if tall:
            A = torch.mm(X, X.mT)
            A2 = torch.mm(A, A)
            B = b * A + c * A2
            torch.mm(B, X, out=bx_buffer)
            X = a * X + bx_buffer
        else:
            A = X @ X.mT
            A2 = A @ A
            B = b * A + c * A2
            X = a * X + B @ X
    if transpose_needed:
        X = X.mT.contiguous() if tall else X.mT
    return X


def out_mm_a2_only(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Reuse only the A@A result buffer; retain exact polynomial expressions."""
    assert G.ndim >= 2
    assert steps <= len(COEFFS)
    transpose_needed = G.size(-2) > G.size(-1)
    tall = transpose_needed and G.ndim == 2
    X = G.bfloat16()
    if tall:
        X = X.mT.contiguous()
    elif transpose_needed:
        X = X.mT
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-7)
    a2_buffer = None
    if tall:
        a2_buffer = torch.empty((X.shape[0], X.shape[0]), device=X.device, dtype=X.dtype)
    for a, b, c in COEFFS[:steps]:
        A = torch.mm(X, X.mT) if tall else X @ X.mT
        if tall:
            torch.mm(A, A, out=a2_buffer)
            A2 = a2_buffer
            B = b * A + c * A2
            X = a * X + torch.mm(B, X)
        else:
            A2 = A @ A
            B = b * A + c * A2
            X = a * X + B @ X
    if transpose_needed:
        X = X.mT.contiguous() if tall else X.mT
    return X


def out_mm_bx_muon_context(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Use the ``out=`` B@X buffer only when autograd cannot be required.

    Muon's optimizer step is decorated with ``torch.no_grad``.  In ordinary
    differentiable calls, this delegates to the exact regular path so the
    public function keeps its autograd contract rather than raising the
    ``out=``/autograd error.
    """
    if torch.is_grad_enabled() and G.requires_grad:
        return contiguous_tall_only(G, steps=steps)
    return out_mm_bx_only(G, steps=steps)


def out_mm_bx_cast_after_context(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Context-safe B@X reuse with cast-after-transpose input preparation."""
    if torch.is_grad_enabled() and G.requires_grad:
        return cast_after_transpose_tall(G, steps=steps)
    assert G.ndim >= 2
    assert steps <= len(COEFFS)
    transpose_needed = G.size(-2) > G.size(-1)
    tall = transpose_needed and G.ndim == 2
    if tall:
        X = G.mT.contiguous().bfloat16()
    else:
        X = G.bfloat16()
        if transpose_needed:
            X = X.mT
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-7)
    bx_buffer = torch.empty_like(X) if tall else None
    for a, b, c in COEFFS[:steps]:
        if tall:
            A = torch.mm(X, X.mT)
            A2 = torch.mm(A, A)
            B = b * A + c * A2
            torch.mm(B, X, out=bx_buffer)
            X = a * X + bx_buffer
        else:
            A = X @ X.mT
            A2 = A @ A
            B = b * A + c * A2
            X = a * X + B @ X
    if transpose_needed:
        X = X.mT.contiguous() if tall else X.mT
    return X


CANDIDATES = {
    "inplace": baseline_order_inplace,
    "addmm": baseline_order_addmm,
    "inplace_add": baseline_order_inplace_add,
    "preallocated_mm": preallocated_mm,
    "contiguous_transpose_mm": contiguous_transpose_mm,
    "contiguous_tall_only": contiguous_tall_only,
    "cast_after_transpose_tall": cast_after_transpose_tall,
    "clone_final_tall": clone_final_tall,
    "out_mm_bx_only": out_mm_bx_only,
    "out_mm_a2_only": out_mm_a2_only,
    "out_mm_bx_muon_context": out_mm_bx_muon_context,
    "out_mm_bx_cast_after_context": out_mm_bx_cast_after_context,
}