def draw_first_k_couples(k: int, rdims: Tensor, dv: torch.device) -> Tensor:
    """Returns first k couples.

    Exhaustive search over the first n samples:
     * n(n+1)/2 = n2/2 + n/2 couples
    Max n for which we can exhaustively sample with k couples:
    * n2/2 + n/2 = k
    * n = sqrt(1/4 + 2k)-1/2 = (sqrt(8k+1)-1)/2
    """
    max_exhaustive_search = int(math.sqrt(2 * k + 0.25) - 0.5)
    residual_search = int(k - max_exhaustive_search * (max_exhaustive_search + 1) / 2)

    repeats = torch.cat(
        [
            torch.arange(max_exhaustive_search, dtype=torch.long, device=dv) + 1,
            torch.tensor([residual_search], dtype=torch.long, device=dv),
        ]
    )
    idx_sequence = torch.stack([repeats.repeat_interleave(repeats), arange_sequence(repeats)], dim=-1)
    return torch.remainder(idx_sequence.unsqueeze(-1), rdims)
Returns first k couples.

    Exhaustive search over the first n samples:
     * n(n+1)/2 = n2/2 + n/2 couples
    Max n for which we can exhaustively sample with k couples:
    * n2/2 + n/2 = k
    * n = sqrt(1/4 + 2k)-1/2 = (sqrt(8k+1)-1)/2
    
kornia/feature/adalam/utils.py