import torch


def cluster_utts(utts: torch.Tensor) -> torch.Tensor:
    """
    given a 2-d tensor of [S][N], where N is number of
    examples, and S is sequence length, and the tensor
    is of discrete int64 indices (cf distributions over
    tokens), we cluster all identical examples, and return
    a cluster assignment as a long tensor, containing the
    cluster id of each example, starting from 0

    if examples have differnet lengths, padding id should
    be identical. this function will compare the entire
    length of each example. as long as the padding id is
    consistent, this should work as desired, i.e. effectively
    ignore padding
    """
    S, N = utts.size()
    clustering = torch.zeros(N, dtype=torch.int64)
    seen = torch.zeros(N, dtype=torch.bool)
    cluster_id = 0
    for n in range(N):
        if seen[n]:
            continue
        mask = (utts == utts[:, n:n + 1]).all(dim=0)
        clustering[mask] = cluster_id
        cluster_id += 1
        seen[mask] = True
    return clustering