code: def compute_batch_indices(batch_size, beam_size):
      """Computes the i'th coordinate that contains the batch index for gathers.
  
      Batch pos is a tensor like [[0,0,0,0,],[1,1,1,1],..]. It says which
      batch the beam item is in. This will create the i of the i,j coordinate
      needed for the gather.
  
      Args:
        batch_size: Batch size
        beam_size: Size of the beam.
      Returns:
        batch_pos: [batch_size, beam_size] tensor of ids
      """
      batch_pos = tf.range(batch_size * beam_size) // beam_size
      batch_pos = tf.reshape(batch_pos, [batch_size, beam_size])
      return batch_pos
  
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 8 for calling a global function with example usage: compute_batch_indices(batch_size, beam_size) and returns: batch_pos