def get_idx_slice(axis, i, ndim):
    """
    Returns a slice `s` where

    mat[s] == mat[:, ---, :, i, :, ---, :]
                  |_______|     |_______|
                  axis          ndim - axis - 1

    This allows indexing into an arbitrary dimension of an n-dimensional array.
    """
    idx = [slice(None)] * ndim
    idx[axis] = i
    return idx