submission 851948

import torch
from task import input_t, output_t

def custom_kernel(data: input_t) -> output_t:
    """
    Optimized Batched Real Symmetric Eigendecomposition.
    Implements a zero-allocation fast path for diagonal matrices
    and leverages UPLO="L" for dense matrices.
    """
    B, N, _ = data.shape
    device = data.device
    dtype = data.dtype
    
    # 1. Fast diagonal check
    nnz_total = torch.count_nonzero(data, dim=(-1, -2))
    nnz_diag = torch.count_nonzero(torch.diagonal(data, dim1=-2, dim2=-1), dim=-1)
    is_diag = (nnz_total == nnz_diag)
    
    # Branch 1: All matrices are perfectly diagonal
    if is_diag.all():
        diag_elements = torch.diagonal(data, dim1=-2, dim2=-1)
        sorted_vals, indices = torch.sort(diag_elements, dim=-1)
        
        Q = torch.zeros((B, N, N), dtype=dtype, device=device)
        b_idx = torch.arange(B, device=device).unsqueeze(1)
        j_idx = torch.arange(N, device=device).unsqueeze(0)
        
        Q[b_idx, indices, j_idx] = 1.0
        return Q, sorted_vals
        
    # Branch 2: All matrices are dense
    elif not is_diag.any():
        # UPLO="L" reduces memory bandwidth by only reading the lower triangle
        values, vectors = torch.linalg.eigh(data, UPLO="L")
        return vectors, values
        
    # Branch 3: Mixed batch
    else:
        values = torch.empty((B, N), dtype=dtype, device=device)
        vectors = torch.empty((B, N, N), dtype=dtype, device=device)
        
        # --- Diagonal Fast Path ---
        diag_idx = is_diag.nonzero(as_tuple=True)[0]
        diag_elements = torch.diagonal(data[diag_idx], dim1=-2, dim2=-1)
        sorted_vals, indices = torch.sort(diag_elements, dim=-1)
        
        # Zero out only the slices for diagonal matrices directly in the output tensor
        # This avoids allocating a temporary (num_diag, N, N) tensor like NIAMEY.py did
        vectors[diag_idx] = 0.0
        
        # Direct assignment of 1.0s to construct the permutation matrices
        b_idx_global = diag_idx.unsqueeze(1)
        j_idx = torch.arange(N, device=device).unsqueeze(0)
        vectors[b_idx_global, indices, j_idx] = 1.0
        
        values[diag_idx] = sorted_vals
        
        # --- Dense Slow Path ---
        not_diag_idx = (~is_diag).nonzero(as_tuple=True)[0]
        
        # Pass only dense matrices to eigh, using UPLO="L" for speed
        slow_values, slow_vectors = torch.linalg.eigh(data[not_diag_idx], UPLO="L")
        
        values[not_diag_idx] = slow_values
        vectors[not_diag_idx] = slow_vectors
        
        return vectors, values