"""Mirrors ``tilelang_ops/asc/gather.py`` (``gather_tl``): bytewise row gather with zero-fill."""

from collections.abc import Callable

import torch


def gather_tl(D: int, dtype: torch.dtype, index_dtype: torch.dtype, x_stride: int) -> Callable[..., None]:
    """Kernel args: ``x [n, D]`` (row stride ``x_stride``), ``index [m]``, ``y [m, D]`` contiguous.

    ``y[i] = x[index[i]]`` if ``0 <= index[i] < n`` else 0; written in place.
    """

    def _gather(x: torch.Tensor, index: torch.Tensor, y: torch.Tensor) -> None:
        n = x.shape[0]
        m = index.shape[0]
        if m == 0:
            return
        if n == 0:
            y.zero_()
            return
        valid = (index >= 0) & (index < n)
        safe = torch.where(valid, index, torch.zeros_like(index))
        rows = torch.index_select(x[:, :D], 0, safe)
        y.copy_(torch.where(valid.view(m, 1), rows, torch.zeros((), dtype=rows.dtype, device=rows.device)))

    return _gather