submission 880494

import torch

from task import input_t, output_t


def _fallback(data: torch.Tensor) -> torch.Tensor:
    return torch.linalg.cholesky_ex(data, check_errors=False).L


def _loop_batch(data: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(data)
    for i in range(data.shape[0]):
        out[i : i + 1].copy_(
            torch.linalg.cholesky_ex(data[i : i + 1], check_errors=False).L
        )
    return out


def custom_kernel(data: input_t) -> output_t:
    n = data.shape[-1]
    batch = data.shape[0]

    # Known win from loop_4096_batch2_experimental.py.
    if n == 4096 and batch == 2:
        return _loop_batch(data)

    # Probe whether other low-batch large shapes have the same bad batched path.
    if n == 2048 and batch in (2, 8):
        return _loop_batch(data)

    if n == 1024 and batch == 4:
        return _loop_batch(data)

    return _fallback(data)