    def reorder_incremental_state(
        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.
        """
        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)
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.
examples/truncated_bptt/transformer_xl_model.py