code: def get_gather_indices(
      X, sequence_lengths, start, end, other_end, chunk_tables
  ):
      """
      Given the tokens, sequence lengths and the range of each token in either rows or columns
      gets the values required for gathering the input tokens into row/column stack inputs and
      scattering the features back into it's original shape.
  
      It also returns the masks to be used by the attention operations and a new set of
      sequence lengths.
  
      Internally - Special tokens are appended to X before the gathers are applied.
  
      Additionally, long rows / columns are chunked up, retaining 2 columns / rows of context (intended to be the headers) into every chunk.
  
      Args:
          X: Input tokens [batch, sequence]
          sequence_lengths: Number of tokens per sequence [batch]
          start: The start of the col/row range.
          end: The end of the col/row range.
          other_end: The end of the row/col range. This is used to determine the first 2 rows / cols to retain in all chunks for context.
          chunk_tables: a boolean. Whether to apply the chunking operation described above
  
      Returns:
      A dict:
      {
          "seq_lens": New sequence lengths [new_batch],
          "values": Indexes into the original sequence of the new batches of data. [new_batch, new_seq_len],
          "attn_mask": Mask for the new sequences [new_batch, new_seq_len, new_seq_len],
      }
  
      """
      seq_len = tf.shape(X)[1]
      bs = tf.shape(X)[0]
  
      mask = tf.expand_dims(
          tf.sequence_mask(sequence_lengths, maxlen=seq_len), 0
      )  # 1, batch, seq
      range = tf.expand_dims(
          tf.expand_dims(
              tf.range(tf.reduce_max(end) + 1, dtype=start.dtype), 1
          ),
          1,
      )  # num_cols, 1, 1
      start = tf.expand_dims(start, 0)  # 1, batch, max_len
      end = tf.expand_dims(end, 0)  # 1, batch, max_len
      col_masks = tf.math.logical_and(
          tf.math.logical_and(
              tf.math.less_equal(start, range),
              tf.math.less_equal(range, end),
          ),
          mask,
      )  # num_cols, batch, seq
  
      batch_idx = tf.tile(tf.expand_dims(tf.range(bs), 1), [1, seq_len])
      seq_idx = tf.tile(tf.expand_dims(tf.range(seq_len), 0), [bs, 1])
      scatter_idx_orig = tf.stack(
          [batch_idx, seq_idx], axis=-1
      )  # batch, seq, 2
      scatter_idx_orig = tf.reshape(scatter_idx_orig, [bs, seq_len, 2])
      return slice_by_table_indices(
          scatter_idx_orig,
          col_masks,
          tf.convert_to_tensor([0, seq_len]),
          tf.convert_to_tensor([0, seq_len + 1]),
          pad_val=tf.convert_to_tensor([0, seq_len + 2]),
          other_end=other_end,
          include_mask=True,
          chunk_tables=chunk_tables,
      )
  
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 1 for calling a global function with example usage: get_gather_indices(X, sequence_lengths, start, end, other_end, chunk_tables) and returns: slice_by_table_indices