def get_diag(arr, i=0):
    """Get the i-th diagonal of a matrix.
    This solution was borrowed from
    http://stackoverflow.com/questions/9958577/changing-the-values-of-the-diagonal-of-a-matrix-in-numpy
    """
    return arr.ravel()[
        max(i, -arr.shape[1] * i) : max(0, (arr.shape[1] - i))
        * arr.shape[1] : arr.shape[1]
        + 1
    ]