import torch


def btrace(M):
    """Computes the trace of batched matrices

    Args:
        M (torch.tensor): matrices of size (Na, Nb, ... Nx, N, N)

    Returns:
        torch.tensor: trace of matrices (Na, Nb, ... Nx)

    Example:
        >>> m = torch.rand(100,5,5)
        >>> tr = btrace(m)
    """
    return torch.diagonal(M, dim1=-2, dim2=-1).sum(-1)