code: def slice_by_table_indices(
      inp,
      col_masks,
      eos_pad,
      bos_pad,
      pad_val,
      other_end,
      base_model_max_length=512,
      max_tokens_per_batch=512 * 100,
      check_len=False,
      include_mask=False,
      chunk_tables=True,
  ):
      """
      Used internally by get_gather_indices - assumes that any pre-processing of adding special tokens is already complete.
      Can handle higher rank inputs but it seems to be quicker to just generate the indices and then gather.
      """
      with tf.device("cpu"):
          bos_pad_ragged = tf.RaggedTensor.from_tensor(
              tf.expand_dims(tf.expand_dims(bos_pad, 0), 0)
          )
          eos_pad_ragged = tf.RaggedTensor.from_tensor(
              tf.expand_dims(tf.expand_dims(eos_pad, 0), 0)
          )
  
          inp_expanded = tf.expand_dims(inp, 0)  # 1, bs, seq, ...
          inp_values_i = tf.ragged.boolean_mask(
              tf.tile(
                  inp_expanded,
                  [
                      tf.shape(col_masks)[0],
                      1,
                      1,
                      *(1 for _ in bos_pad.shape),
                  ],
              ),
              col_masks,
          ).merge_dims(0, 1)
          batch_mask = tf.math.not_equal(inp_values_i.row_lengths(), 0)
          inp_values = tf.RaggedTensor.from_row_lengths(
              values=inp_values_i.flat_values,
              row_lengths=tf.boolean_mask(
                  inp_values_i.row_lengths(), batch_mask
              ),
          )
          if chunk_tables:
              inp_values = chunk_ragged_tensor(
                  inp_values,
                  other_end=other_end,
                  base_model_max_length=base_model_max_length - 2,
              )  # -2 for eos and bos
          col_bs = tf.shape(inp_values.row_lengths())[0]
          bos_expanded = tf.tile(
              bos_pad_ragged, [col_bs, 1, *(1 for _ in bos_pad.shape)]
          )
          eos_expanded = tf.tile(
              eos_pad_ragged, [col_bs, 1, *(1 for _ in eos_pad.shape)]
          )
          output_ragged = tf.concat(
              [bos_expanded, inp_values, eos_expanded], axis=1
          )
  
          output_ragged, mask, pos_ids = batch_packing(
              output_ragged,
              include_mask=include_mask,
              base_model_max_length=base_model_max_length,
          )
          col_seq_lens = output_ragged.row_lengths()
  
          # I don't think this default matters here as these are masked. Cannot be < 0
          col_values = output_ragged.to_tensor(
              default_value=pad_val, shape=[None, None] + inp.shape[2:]
          )
          pos_ids = pos_ids.to_tensor(
              default_value=0, shape=[None, None]
          )
      # This is to handle some edge cases where 0 length values are missing the final dim after conversion.
      col_values.set_shape([None, None] + list(inp.shape[2:]))
      max_length = tf.minimum(
          tf.math.floordiv(
              max_tokens_per_batch, tf.shape(col_values)[0]
          ),
          base_model_max_length,
      )
      if check_len:
          ctrl_dep = [
              tf.cond(
                  tf.shape(col_values)[1] > max_length,
                  false_fn=tf.no_op,
                  true_fn=lambda: tf.compat.v1.Print(
                      True,
                      [
                          "The length of produced tensors is >",
                          max_length,
                          "shape is: ",
                          tf.shape(col_values),
                      ],
                  ),
              )
          ]
      else:
          ctrl_dep = []
      if not chunk_tables:
          with tf.control_dependencies(ctrl_dep):
              col_values = col_values[:, :max_length]
              if include_mask:
                  mask = mask[:, :max_length, :max_length]
                  mask.set_shape([None, None, None])
              pos_ids = pos_ids[:, :max_length]
  
      col_values.set_shape([None, None] + list(inp.shape[2:]))
      return {
          "seq_lens": col_seq_lens,
          "values": col_values,
          "attn_mask": mask,
          "pos_ids": pos_ids,
      }
  
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 4 for calling a global function with example usage: slice_by_table_indices(inp, col_masks, eos_pad, bos_pad, pad_val, other_end, base_model_max_length, max_tokens_per_batch, check_len, include_mask, chunk_tables) and returns: dict