import torch


def batched_select(data, indices):
    """
    Select data[i, indices[i]] for each batch element i.

    Args:
        data: [B, N, ...] input tensor
        indices: [B, K] indices for each batch element

    """

    assert data.shape[0] == indices.shape[0], f"Batch size mismatch {data.shape[0]} vs {indices.shape[0]}"
    assert indices.dim() == 2, f"indices should be 2D, got {indices.shape}"

    B = data.shape[0]
    batch_idx = torch.arange(B, device=data.device)[:, None]
    return data[batch_idx, indices]