code: def batch_packing(
      ragged_input, include_mask=True, base_model_max_length=512
  ):
      """
      Takes a ragged tensor input and re-packs the batches to minimise the batch size without
      impacting the sequence length. Additionally returns a mask to use with self-attention so
      that there is no change in output.
  
      Args:
          ragged_input: A ragged tensor of at least rank 2. - The inputs to repack.
          include_mask (bool, optional): Whether to return the masks.
  
      Returns:
          new output - repacked but still ragged.
          the mask [batch, seq_len, seq_len]
          position ids [batch, seq_len]
      """
      # Basically the next fit binpacking algorithm.
  
      # If there is some fast way to sort the rows of a ragged tensor we might consider doing this first.
      col_seq_lens = ragged_input.row_lengths()
      max_length = tf.minimum(
          tf.reduce_max(col_seq_lens), base_model_max_length
      )
  
      def loop_body(l, i):
          temp_ragged = tf.RaggedTensor.from_row_lengths(
              values=col_seq_lens[:i], row_lengths=l
          )
          return (
              tf.cond(
                  pred=tf.math.less_equal(
                      tf.minimum(
                          tf.reduce_sum(temp_ragged[-1]),
                          base_model_max_length,
                      ),
                      max_length - col_seq_lens[i],
                  ),
                  true_fn=lambda: tf.concat(
                      (l[:-1], [l[-1] + 1]), axis=0
                  ),
                  false_fn=lambda: tf.concat((l, [1]), axis=0),
              ),
              i + 1,
          )
  
      lengths, _ = tf.while_loop(
          cond=lambda _, i: tf.math.less(i, tf.shape(col_seq_lens)[0]),
          body=loop_body,
          loop_vars=(tf.ones([1], dtype=tf.int64), tf.constant(1)),
          parallel_iterations=1,
          shape_invariants=(tf.TensorShape([None]), tf.TensorShape([])),
      )
  
      lengths = tf.RaggedTensor.from_row_lengths(
          values=col_seq_lens, row_lengths=lengths
      )
      ragged = tf.RaggedTensor.from_row_lengths(
          values=ragged_input.flat_values,
          row_lengths=tf.reduce_sum(lengths, 1),
      )
      if include_mask:
          lengths_ts = tf.math.cumsum(lengths.to_tensor(), 1)
          raw_mask = tf.sequence_mask(
              lengths_ts,
              dtype=tf.float32,
              maxlen=tf.minimum(
                  base_model_max_length,
                  tf.cast(tf.reduce_max(lengths_ts), tf.int32),
              ),
          )
          compound_mask = tf.reduce_sum(
              tf.expand_dims(raw_mask, 2) * tf.expand_dims(raw_mask, 3),
              1,
          )
          row_max = tf.reduce_max(compound_mask, 1)
          output_mask = tf.cast(
              tf.logical_and(
                  tf.equal(compound_mask, tf.expand_dims(row_max, 1)),
                  tf.equal(compound_mask, tf.expand_dims(row_max, 2)),
              ),
              tf.float32,
          )
  
      else:
          output_mask = None
  
      pos_ids = tf.RaggedTensor.from_row_lengths(
          values=tf.ragged.range(lengths.flat_values).flat_values,
          row_lengths=tf.reduce_sum(lengths, 1),
      )
      return ragged, output_mask, 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 2 for calling a global function with example usage: batch_packing(ragged_input, include_mask, base_model_max_length) and returns: ragged,  output_mask,  pos_ids, tf,  unknown