def off_diagonal(x):
    """ Return a flattened view of the off-diagonal elements of a square
    matrix.
    """
    n, m = x.shape
    assert n == m
    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()