def METHOD_NAME(
    self,
    incremental_state: Dict[str, Dict[str, Optional[Tensor]]],
    new_order: 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.
    """
    pass