class_docstr: A `Sampler` which redistributes the probability mass function for nodes among the
  top `k` choices, then samples from that subset after re-normalizing the probabilities.
  
  Beams are sampled in the default, deterministic way.
  
  # Parameters
  
  k : `int`, optional (default = `1`)
      The number of top choices to be selected from.
  temperature : `float`, optional (default = `1.0`)
      A `temperature` below 1.0 produces a sharper probability distribution and a `temperature`
      above 1.0 produces a flatter probability distribution.
  with_replacement: `bool`, optional, (default = `False`)
      If set to `True`, samples will be selected with replacement from the top k choices.
class_docstr_tok: ["A","`","Sampler","`","which","redistributes","the","probability","mass","function","for","nodes","among","the","top","`","k","`","choices",",","then","samples","from","that","subset","after","re-normalizing","the","probabilities",".","Beams","are","sampled","in","the","default",",","deterministic","way",".","#","Parameters","k",":","`","int","`",",","optional","(","default","=","`","1","`",")","The","number","of","top","choices","to","be","selected","from",".","temperature",":","`","float","`",",","optional","(","default","=","`","1.0","`",")","A","`","temperature","`","below","1.0","produces","a","sharper","probability","distribution","and","a","`","temperature","`","above","1.0","produces","a","flatter","probability","distribution",".","with_replacement",":","`","bool","`",",","optional",",","(","default","=","`","False","`",")","If","set","to","`","True","`",",","samples","will","be","selected","with","replacement","from","the","top","k","choices","."]
code: def sample_nodes(
      self,
      log_probs: torch.Tensor,
      per_node_beam_size: int,
      state: StateType,
  ) -> Tuple[torch.Tensor, torch.Tensor, StateType]:
      if not per_node_beam_size <= self.k <= log_probs.size()[1]:
          raise ValueError(
              "k must be a postive integer no less than per_node_beam_size and no greater than vocabulary size"
          )
  
      # shape (both): (batch_size, k)
      top_k_log_probs, top_k_indices = log_probs.topk(self.k, dim=-1)
  
      # Apply temperature if necessary.
      # shape: (batch_size, k)
      if self.temperature != 1.0:
          top_k_log_probs = top_k_log_probs / self.temperature
  
      # Re-normalize the subset.
      # shape: (batch_size, k)
      normalized_top_k_probs = torch.nn.functional.softmax(
          top_k_log_probs, dim=-1
      )
  
      # Sample from the re-normalized subset.
      # NOTE: These indices are not indices into `log_probs`, they are indices into `top_k_log_probs`.
      # shape: (batch_size, per_node_beam_size)
      sampled_indices = torch.multinomial(
          normalized_top_k_probs,
          per_node_beam_size,
          replacement=self.with_replacement,
      )
  
      # Convert `sampled_indices` back to indices in the original `log_probs` tensor.
      # shape: (batch_size, per_node_beam_size)
      indices = top_k_indices.gather(-1, sampled_indices)
  
      return log_probs.gather(1, indices), indices, state
  
imports: from inspect import signature
  from typing import Any, List, Callable, Tuple, Dict, cast, TypeVar, Optional
  import copy
  import warnings
  import torch
  from allennlp.common import Lazy, Registrable
  from allennlp.common.checks import ConfigurationError
  from allennlp.data import Vocabulary
  from allennlp.nn.util import min_value_of_dtype
label_desc: Use image node_id 2 for calling the TopKSampler obj's underlying member method code with example usage: obj.sample_nodes(log_probs, per_node_beam_size, state) and returns: log_probs,  indices,  state