def diagonal(mat, diag_index):
    """Returns ith diagonal of matrix, where i is the diag_index.

    Returns the ith diagonal (A_0i, A_1(i+1), ..., A_N(i-1)) of a matrix A,
    where i is the diag_index.

    Args:
        mat (2-D list): Matrix.
        diag_index (int): Index of diagonal to return.

    Returns:
        Diagonal of a matrix.
    """
    return [mat[j % len(mat)][(diag_index + j) % len(mat)] for j in range(len(mat))]