code: def gather_col_vals(inp, gather_output, eos_pad, bos_pad, pad_val):
      """Used with the output of get_gather_indices to gather the input to the models.
  
      Args:
          inp: input to be gathered at least rank 2 with the first 2 dimensions [batch, sequence]_
          gather_output: THe output of get_gather_indices
          eos_pad: A value to be used as the end of sequence value - equal to the rank of inp - 2.
          bos_pad: A value to be used as the beginning of sequence value - equal to the rank of inp - 2.
          pad_val A value to be used as the padding value - equal to the rank of inp - 2.
  
      Returns:
          a dictionary
          {
              "seq_lens": seq_lens from gather output unmodified,
              "values": the values inp rearraged to be ready as input to the model,
              "attn_mask": mask from gather output unmodified,
          }
      """
      bs = tf.shape(inp)[0]
      bos_expanded = tf.tile(
          tf.expand_dims(tf.expand_dims(bos_pad, 0), 0),
          [bs, 1, *(1 for _ in bos_pad.shape)],
      )
      eos_expanded = tf.tile(
          tf.expand_dims(tf.expand_dims(eos_pad, 0), 0),
          [bs, 1, *(1 for _ in eos_pad.shape)],
      )
      inp_w_extra_toks = tf.concat(
          [
              inp,
              bos_expanded,
              eos_expanded,
              tf.ones_like(eos_expanded) * pad_val,
          ],
          axis=1,
      )
      values = tf.gather_nd(
          indices=gather_output["values"], params=inp_w_extra_toks
      )
      return {
          "seq_lens": gather_output["seq_lens"],
          "values": values,
          "attn_mask": gather_output["attn_mask"],
      }
  
file_path: finetune/finetune/base_models/bert/table_utils.py
imports: import tensorflow
  from finetune.base_models.bert.modeling import gelu
label_desc: Use image node_id 5 for calling a global function with example usage: gather_col_vals(inp, gather_output, eos_pad, bos_pad, pad_val) and returns: dict