import torch


def trace(mat):
    """Return the N traces of a batch of N square matrices,
    or return the trace of a square matrix."""
    # Default batch size is 1
    if mat.dim() < 3:
        mat = mat.unsqueeze(dim=0)
    if mat.is_cuda:
        tr =  (torch.eye(mat.shape[1], dtype=mat.dtype).cuda() * mat).sum(dim=1).sum(dim=1)
    # Element-wise multiply by identity and take the sum
    else:
        tr =  (torch.eye(mat.shape[1], dtype=mat.dtype) * mat).sum(dim=1).sum(dim=1)
    
    return tr.view(mat.shape[0])