def apply_mapping(mask, mapping):
    """Applies the mapping to obtain a frequency aligned mask.

    Args:
        mask: Permuted mask with shape (K, F, ...).
        mapping: Reverse mapping with shape (K, F).

    >>> np.random.seed(0)
    >>> K, F, T = 3, 5, 6
    >>> reference_mask = np.zeros((K, F, T), dtype=np.int8)
    >>> def viz_mask(mask: np.ndarray):
    ...     mask = np.einsum('KFT->FKT', mask).astype(str).tolist()
    ...     for mask_f in mask:
    ...         print('   '.join([' '.join(m) for m in mask_f]))
    >>> reference_mask[0, :, :2] = 1
    >>> reference_mask[1, :, 2:4] = 1
    >>> reference_mask[2, :, 4:] = 1
    >>> viz_mask(reference_mask)
    1 1 0 0 0 0   0 0 1 1 0 0   0 0 0 0 1 1
    1 1 0 0 0 0   0 0 1 1 0 0   0 0 0 0 1 1
    1 1 0 0 0 0   0 0 1 1 0 0   0 0 0 0 1 1
    1 1 0 0 0 0   0 0 1 1 0 0   0 0 0 0 1 1
    1 1 0 0 0 0   0 0 1 1 0 0   0 0 0 0 1 1
    >>> mapping = sample_random_mapping(K, F)
    >>> mask = apply_mapping(reference_mask, mapping)
    >>> viz_mask(mask)
    0 0 0 0 1 1   0 0 1 1 0 0   1 1 0 0 0 0
    0 0 0 0 1 1   1 1 0 0 0 0   0 0 1 1 0 0
    1 1 0 0 0 0   0 0 0 0 1 1   0 0 1 1 0 0
    0 0 0 0 1 1   1 1 0 0 0 0   0 0 1 1 0 0
    0 0 0 0 1 1   0 0 1 1 0 0   1 1 0 0 0 0

    Test against a loopy implementation of apply mapping
    >>> def apply_mapping_loopy(mask, mapping):
    ...     _, F = mapping.shape
    ...     aligned_mask = np.zeros_like(mask)
    ...     for f in range(F):
    ...         aligned_mask[:, f, :] = mask[mapping[:, f], f, :]
    ...     return aligned_mask
    >>> mask = apply_mapping_loopy(reference_mask, mapping)
    >>> viz_mask(mask)
    0 0 0 0 1 1   0 0 1 1 0 0   1 1 0 0 0 0
    0 0 0 0 1 1   1 1 0 0 0 0   0 0 1 1 0 0
    1 1 0 0 0 0   0 0 0 0 1 1   0 0 1 1 0 0
    0 0 0 0 1 1   1 1 0 0 0 0   0 0 1 1 0 0
    0 0 0 0 1 1   0 0 1 1 0 0   1 1 0 0 0 0
    """
    K, F = mapping.shape
    assert K < 20, (K, mapping.shape)
    assert mask.shape[:2] == mapping.shape, (mask.shape, mapping.shape)
    return mask[mapping, range(F)]