submission 864397

#!POPCORN leaderboard eigh
#!POPCORN gpu B200
# ---
# parent: none
# intent: baseline torch.linalg.eigh from the scaffold
# status: pending
# result: null
# id: null
# ---
"""Baseline batched symmetric eigendecomposition via torch.linalg.eigh."""

import torch
from task import input_t, output_t

# Contract gotchas
# Input is one batch x n x n FP32 CUDA tensor, symmetric up to FP32 roundoff.
# Return order is (Q, L), the reverse of torch.linalg.eigh which gives
# (values, vectors). L must be ascending, both outputs FP32 on the input
# device and finite, NaN or Inf anywhere is an immediate fail.
# The checker never compares against reference eigenvectors. It gates on
# A @ Q - Q @ diag(L), reconstruction Q @ diag(L) @ Q.T - A, orthogonality
# Q.T @ Q - I, each relative to matrix L1 norms in FP64, plus a tolerant
# ascending sort check. Sign flips and rotations inside eigenspaces are fine.
# Stress inputs include zero matrices, rank deficient, repeated and clustered
# spectra, and magnitudes near sqrt(float32 max) and sqrt(float32 tiny).


def custom_kernel(data: input_t) -> output_t:
    """Solve batched eigh and return eigenvectors then eigenvalues."""
    values, vectors = torch.linalg.eigh(data)
    return vectors, values  # pyright: ignore[reportReturnType]