def METHOD_NAME(
    self,
    incremental_state: Dict[str, Dict[str, Optional[torch.Tensor]]],
    new_order: torch.Tensor,
):
    """Reorder incremental state.
    This will be called when the order of the input has changed from the
    previous time step. A typical use case is beam search, where the input
    order changes between time steps based on the selection of beams.
    """
    raise NotImplementedError("This is required for generation/beam search")
    # mems = self.get_incremental_state(incremental_state, "mems")
    # if mems is not None:
    #     new_mems = [mems_i.index_select(1, new_order) for mems_i in mems]
    #     self.set_incremental_state(incremental_state, "mems", new_mems)