code: def sort_state(self, best_hyp_indices: mx.nd.NDArray):
          
          self.states = [mx.nd.take(ds, best_hyp_indices) for ds in self.states]
docs: Sorts states according to k-best order from last step in beam search.