Compare eager and compiled normalization on real tensors

"""Validate recorded numerical parity and summarize actual Metal expression timings."""

import json, math, statistics
from pathlib import Path


def compare_case(case, tolerance=1e-5):
    if (
        tolerance < 0
        or not math.isfinite(case["max_abs_error"])
        or case["max_abs_error"] > tolerance
    ):
        raise ValueError("Numerical parity exceeds tolerance")
    result = {
        "rows": case["rows"],
        "width": case["width"],
        "max_abs_error": case["max_abs_error"],
    }
    for name in ["eager", "compiled"]:
        samples = case["timings"][name]["samples_ms"]
        if not samples or any(not math.isfinite(v) or v <= 0 for v in samples):
            raise ValueError("Invalid timing samples")
        result[name + "_median_ms"] = round(statistics.median(samples), 6)
    return result


def main():
    record = json.loads(Path("fusion-run.json").read_text())
    print(
        json.dumps(
            {
                "execution_mode": "offline analysis of fresh Metal eager and compiled expression runs",
                "cases": [compare_case(c) for c in record["cases"]],
                "limits": "Shape-specific tensor timings; no Triton/CUDA run, device-memory counter or whole-model speedup.",
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()

"""Compare eager and compiled RMS normalization on Metal; no CUDA/Triton execution claimed."""

import argparse
import hashlib
import importlib.metadata
import json
import platform
import statistics
import time
from pathlib import Path


def main():
    import mlx.core as mx

    p = argparse.ArgumentParser()
    p.add_argument("--output", required=True)
    a = p.parse_args()
    mx.random.seed(0)

    def norm(x, r, w):
        z = x + r
        return z * mx.rsqrt(mx.mean(z * z, axis=-1, keepdims=True) + 1e-6) * w

    compiled = mx.compile(norm)
    cases = []
    for rows, width in [(1, 896), (32, 896), (32, 1003)]:
        x = mx.random.normal((rows, width))
        r = mx.random.normal((rows, width))
        w = mx.ones((width,))
        mx.eval(x, r, w)
        reference = norm(x, r, w)
        candidate = compiled(x, r, w)
        mx.eval(reference, candidate)
        error = float(mx.max(mx.abs(reference - candidate)).item())
        timings = {}
        for name, fn in [("eager", norm), ("compiled", compiled)]:
            for _ in range(3):
                mx.eval(fn(x, r, w))
            samples = []
            for _ in range(20):
                mx.synchronize()
                start = time.perf_counter()
                out = fn(x, r, w)
                mx.eval(out)
                mx.synchronize()
                samples.append((time.perf_counter() - start) * 1000)
            timings[name] = {
                "median_ms": round(statistics.median(samples), 6),
                "samples_ms": [round(s, 6) for s in samples],
            }
        cases.append(
            {
                "rows": rows,
                "width": width,
                "dtype": "float32",
                "max_abs_error": error,
                "timings": timings,
            }
        )
    result = {
        "execution_mode": "fresh Metal tensor execution; eager versus MLX compiled expression",
        "environment": {
            "python": platform.python_version(),
            "platform": platform.platform(),
            "mlx": importlib.metadata.version("mlx"),
        },
        "source_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        "cases": cases,
        "limits": "Expression compilation may fuse operations. No kernel trace or measured HBM byte count; no whole-model, CUDA or Triton speedup claim.",
    }
    Path(a.output).write_text(json.dumps(result, indent=2) + "\n")


if __name__ == "__main__":
    main()

import unittest
from lab import compare_case


class FusionEvidenceTests(unittest.TestCase):
    def case(self):
        return {
            "rows": 1,
            "width": 3,
            "max_abs_error": 0,
            "timings": {
                "eager": {"samples_ms": [1, 3, 2]},
                "compiled": {"samples_ms": [2, 2, 2]},
            },
        }

    def test_parity_is_required_before_timing(self):
        case = self.case()
        case["max_abs_error"] = 0.1
        with self.assertRaises(ValueError):
            compare_case(case)

    def test_invalid_execution_timing_is_rejected(self):
        case = self.case()
        case["timings"]["compiled"]["samples_ms"] = [-1]
        with self.assertRaises(ValueError):
            compare_case(case)

    def test_medians_are_compared_without_a_speedup_promise(self):
        result = compare_case(self.case())
        self.assertEqual(result["eager_median_ms"], 2)
        self.assertEqual(result["compiled_median_ms"], 2)


if __name__ == "__main__":
    unittest.main()