import torch
from typing import Tuple


def do_rnnt_pruning(
    am: torch.Tensor, lm: torch.Tensor, ranges: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Prune the output of encoder(am) output and prediction network(lm)
    output of RNNT.

    Args:
      am:
        The encoder output, with shape (B, T, C)
      lm:
        The prediction network output, with shape (B, S + 1, C)
      ranges:
        A tensor containing the symbol indexes for each frame that we want to
        keep. Its shape is (B, T, s_range), see the docs in
        `get_rnnt_prune_ranges` for more details of this tensor.

    Returns:
      Return the pruned am and lm with shape (B, T, s_range, C)
    """
    # am (B, T, C)
    # lm (B, S + 1, C)
    # ranges (B, T, s_range)
    assert ranges.shape[0] == am.shape[0]
    assert ranges.shape[0] == lm.shape[0]
    assert am.shape[1] == ranges.shape[1]
    (B, T, s_range) = ranges.shape
    (B, S1, C) = lm.shape
    S = S1 - 1

    # (B, T, s_range, C)
    am_pruning = am.unsqueeze(2).expand((B, T, s_range, C))

    # (B, T, s_range, C)
    lm_pruning = torch.gather(
        lm.unsqueeze(1).expand((B, T, S + 1, C)),
        dim=2,
        index=ranges.reshape((B, T, s_range, 1)).expand((B, T, s_range, C)),
    )
    return am_pruning, lm_pruning