nodes: 0, module; 1, function_definition; 2, function_name:sort_state; 3, parameters; 4, block; 5, identifier:self; 6, typed_parameter; 7, expression_statement; 8, expression_statement; 9, identifier:best_hyp_indices; 10, type; 11, comment:"""
          Sorts states according to k-best order from last step in beam search.
          """; 12, assignment; 13, attribute; 14, attribute; 15, list_comprehension; 16, attribute; 17, identifier:NDArray; 18, identifier:self; 19, identifier:states; 20, call; 21, for_in_clause; 22, identifier:mx; 23, identifier:nd; 24, attribute; 25, argument_list; 26, identifier:ds; 27, attribute; 28, attribute; 29, identifier:take; 30, identifier:ds; 31, identifier:best_hyp_indices; 32, identifier:self; 33, identifier:states; 34, identifier:mx; 35, identifier:nd
code: def sort_state(self, best_hyp_indices: mx.nd.NDArray):
          """
          Sorts states according to k-best order from last step in beam search.
          """
          self.states = [mx.nd.take(ds, best_hyp_indices) for ds in self.states]