code: def soft_top_k(flat_logits, temperature=1.0, sample_from=1, k=1):
      flat_logits /= tf.cast(temperature, dtype=tf.float32)
      flat_logits = top_k_logits(flat_logits, k=sample_from)
      return tf.random.categorical(
          flat_logits, num_samples=k, dtype=tf.int32
      )
  
file_path: finetune/finetune/util/beam_search.py
imports: import math
  import numpy
  from finetune.util.shapes import shape_list
  import tensorflow
  from tensorflow.python.ops import inplace_ops
  from tensorflow.python.util import nest
label_desc: Use image node_id 2 for calling a global function with example usage: soft_top_k(flat_logits, temperature, sample_from, k) and returns: tf