code: def chunk_ragged_tensor(
      inputs: tf.RaggedTensor,
      other_end: tf.Tensor,
      include_n_rows: int = 2,
      base_model_max_length: int = 512 - 2,
      # -2 to account for BOS and EOS tokens that will be added.
  ):
      """
      Chunks up inputs, keeping the first include_n_rows in every chunk.
      """
      batch_size = inputs.bounding_shape()[0]
  
      # Figure out which tokens are in the first include_n_rows rows of the column and split these off.
  
      is_first_n_rows_orig = tf.gather_nd(
          tf.math.less(other_end, include_n_rows), inputs, batch_dims=0
      )
      less_than_max_length_tokens = tf.math.less(
          tf.math.reduce_sum(
              tf.cast(is_first_n_rows_orig, tf.int64), axis=1
          ),
          base_model_max_length,
      )
      # If first row tokens is more than base_model_max_length - do not use any first row context.
      is_first_n_rows = tf.logical_and(
          tf.expand_dims(less_than_max_length_tokens, 1),
          is_first_n_rows_orig,
      )
      num_first_n_row_tokens = tf.math.reduce_sum(
          tf.cast(is_first_n_rows, tf.int64), axis=1
      )
  
      num_cols_per_col = tf.cast(
          tf.maximum(
              tf.math.ceil(
                  (inputs.row_lengths() - num_first_n_row_tokens)
                  / (base_model_max_length - num_first_n_row_tokens)
              ),
              1,  # Max(1 .) is to catch the case where all tokens are in the first n rows.
          ),
          tf.int64,
      )
      first_n_rows = tf.ragged.boolean_mask(inputs, is_first_n_rows)
      other_rows = tf.ragged.boolean_mask(
          inputs, tf.math.logical_not(is_first_n_rows)
      )
      other_rows_length_goals = (
          base_model_max_length - num_first_n_row_tokens
      )  # bs
  
      # Reshape the rest of the columns to the per-column length_goal
      new_row_breaks = tf.cast(
          tf.minimum(
              tf.ragged.range(
                  (num_cols_per_col + 1) * other_rows_length_goals,
                  deltas=other_rows_length_goals,
                  dtype=tf.int64,
              ),
              tf.expand_dims(other_rows.row_lengths(), 1),
          ),
          tf.int64,
      )
      new_row_lengths = new_row_breaks[:, 1:] - new_row_breaks[:, :-1]
      other_rows_reshaped = tf.RaggedTensor.from_row_lengths(
          values=other_rows.flat_values,
          row_lengths=new_row_lengths.flat_values,
      )
  
      # Tile the first N rows so that they are duplicated for each chunk of the column
  
      # a really gross way to get range broadcasted to num_cols_per_col - not sure of a better way to get this.
      indexes = tf.ragged.range(num_cols_per_col) * 0 + tf.expand_dims(
          tf.range(batch_size), 1
      )
      first_n_rows_duped = tf.gather(first_n_rows, indexes).merge_dims(
          0, 1
      )
  
      result = tf.concat([first_n_rows_duped, other_rows_reshaped], 1)
      return result
  
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 3 for calling a global function with example usage: chunk_ragged_tensor(inputs, other_end, include_n_rows, base_model_max_length) and returns: result