submission 874653

#!POPCORN leaderboard eigh
#!POPCORN gpu B200

import torch
from task import input_t, output_t


def _diagonal_eigh(data: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    # For a diagonal matrix the answer is free:
    #   eigenvalues  = the diagonal entries (sorted ascending)
    #   eigenvectors = the standard basis vectors e_i, reordered the same way
    diag = data.diagonal(dim1=-2, dim2=-1)          # batch x n
    values, order = torch.sort(diag, dim=-1)        # ascending, per matrix

    # Column k of Q must be the basis vector for the k-th smallest eigenvalue.
    # one_hot(order) puts a 1 at [b, k, order[b, k]]; transposing makes it column k.
    vectors = torch.nn.functional.one_hot(order, num_classes=data.shape[-1])
    vectors = vectors.transpose(-2, -1).to(data.dtype)
    return vectors, values


def custom_kernel(data: input_t) -> output_t:
    # Fast path: if every off-diagonal entry is exactly zero, skip the solver.
    off_diag = data - torch.diag_embed(data.diagonal(dim1=-2, dim2=-1))
    if off_diag.abs().amax().item() == 0.0:
        return _diagonal_eigh(data)

    # General case: same solver as the baseline.
    values, vectors = torch.linalg.eigh(data)
    return vectors, values