import torch


def matperm2listperm(matperm):
    """Converts permutation matrix to its enumeration (list) form.

    Args:
    matperm: (..., n, n)

    Returns:
    listperm: (..., n) - listperm[t,i] is the index of the only non-zero entry in matperm[t, i, :]
    """
    batch_size = matperm.size(0)
    n = matperm.size(-1)
    assert matperm.size(-2) == matperm.size(-1)

    #argmax is the index location of each maximum value found(argmax)
    # _, argmax = torch.max(matperm, dim=-1, keepdim=False)
    argmax = torch.argmax(matperm, dim=-1, keepdim=False)
    # argmax = argmax.view(batch_size, n_objects)
    return argmax