code: def consume(self, word_ids: mx.nd.NDArray) -> None:
          
          word_ids = word_ids.asnumpy().tolist()
          for i, word_id in enumerate(word_ids):
              if self.global_avoid_states:
                  self.global_avoid_states[i] = self.global_avoid_states[i].consume(word_id)
              if self.local_avoid_states:
                  self.local_avoid_states[i] = self.local_avoid_states[i].consume(word_id)
docs: Consumes a word for each trie, updating respective states.
  
          :param word_ids: The set of word IDs.