submission 839293

import torch

from task import input_t, output_t


def _identity(batch: int, n: int, data: torch.Tensor) -> torch.Tensor:
    return torch.eye(n, device=data.device, dtype=torch.float32).expand(batch, n, n)


def _diagonal_eigh(data: torch.Tensor) -> output_t | None:
    batch, n, _ = data.shape
    if batch > 2:
        return None
    if n > 1 and bool(data[..., 0, 1].abs().amax() != 0):
        return None
    diag = torch.diagonal(data, dim1=-2, dim2=-1)
    offdiag_l1 = data.abs().sum(dim=(-2, -1)) - diag.abs().sum(dim=-1)
    if bool(offdiag_l1.amax() != 0):
        return None
    values, order = torch.sort(diag, dim=-1)
    vectors = _identity(batch, n, data).gather(2, order.unsqueeze(1).expand(batch, n, n))
    return vectors.contiguous(), values.contiguous()


def custom_kernel(data: input_t) -> output_t:
    diagonal = _diagonal_eigh(data)
    if diagonal is not None:
        return diagonal
    values, vectors = torch.linalg.eigh(data)
    return vectors, values