anima_mlx/utils/compare.py

"""Numeric comparison helpers for golden parity tests."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any


@dataclass(frozen=True)
class TensorComparison:
    shape: tuple[int, ...]
    reference_shape: tuple[int, ...]
    max_abs: float
    mean_abs: float
    rmse: float
    max_rel: float
    passed: bool


def to_numpy(array: Any) -> Any:
    """Convert PyTorch/MLX/NumPy-like arrays to a NumPy float32 array."""

    import numpy as np

    if hasattr(array, "detach"):
        return array.detach().cpu().float().numpy()
    if hasattr(array, "tolist"):
        return np.asarray(array.tolist(), dtype=np.float32)
    return np.asarray(array, dtype=np.float32)


def compare_tensors(
    actual: Any,
    expected: Any,
    *,
    atol: float,
    rtol: float = 0.0,
) -> TensorComparison:
    """Return absolute/relative error metrics and threshold result."""

    import numpy as np

    actual_np = to_numpy(actual).astype(np.float32)
    expected_np = to_numpy(expected).astype(np.float32)
    if actual_np.shape != expected_np.shape:
        return TensorComparison(
            shape=tuple(actual_np.shape),
            reference_shape=tuple(expected_np.shape),
            max_abs=float("inf"),
            mean_abs=float("inf"),
            rmse=float("inf"),
            max_rel=float("inf"),
            passed=False,
        )

    diff = np.abs(actual_np - expected_np)
    denom = np.maximum(np.abs(expected_np), 1e-12)
    rel = diff / denom
    allowed = atol + rtol * np.abs(expected_np)
    passed = bool(np.all(diff <= allowed))
    return TensorComparison(
        shape=tuple(actual_np.shape),
        reference_shape=tuple(expected_np.shape),
        max_abs=float(diff.max()) if diff.size else 0.0,
        mean_abs=float(diff.mean()) if diff.size else 0.0,
        rmse=float(np.sqrt(np.mean(np.square(diff)))) if diff.size else 0.0,
        max_rel=float(rel.max()) if rel.size else 0.0,
        passed=passed,
    )