submission 838388

#!POPCORN leaderboard eigh
#!POPCORN gpu B200

# Batched real-symmetric eigendecomposition.
#
# Baseline: torch.linalg.eigh dispatches to cuSOLVER's batched Jacobi (syevj),
# which is well-optimized at every size we hit (B200 profiling confirmed it uses
# batch_parallel_jacobi for n=32 and syevj for n=512). A naive custom Jacobi was
# measured ~5x slower at n=32, so there is no cheap win in the small cases.
# Beating this requires a fundamentally faster algorithm (e.g. a low-precision
# tensor-core spectral solver exploiting the loose residual gates).

import torch
from task import input_t, output_t


def custom_kernel(data: input_t) -> output_t:
    values, vectors = torch.linalg.eigh(data)
    return vectors, values