def collapse_dims(rt, axis=0):
  """Collapses the specified axis of a RaggedTensor.

  Suppose we have a RaggedTensor like this:
  [[1, 2, 3],
   [4, 5],
   [6]]

  If we flatten the 0th dimension, it becomes:
  [1, 2, 3, 4, 5, 6]

  Args:
    rt: a RaggedTensor.
    axis: the dimension to flatten.

  Returns:
    A flattened RaggedTensor, which now has one less dimension.
  """
  to_expand = rt.nested_row_lengths()[axis]
  to_elim = rt.nested_row_lengths()[axis + 1]

  bar = tf.RaggedTensor.from_row_lengths(to_elim, row_lengths=to_expand)
  new_row_lengths = tf.reduce_sum(bar, axis=1)
  return tf.RaggedTensor.from_nested_row_lengths(
      rt.flat_values,
      rt.nested_row_lengths()[:axis] + (new_row_lengths,))
Collapses the specified axis of a RaggedTensor.

  Suppose we have a RaggedTensor like this:
  [[1, 2, 3],
   [4, 5],
   [6]]

  If we flatten the 0th dimension, it becomes:
  [1, 2, 3, 4, 5, 6]

  Args:
    rt: a RaggedTensor.
    axis: the dimension to flatten.

  Returns:
    A flattened RaggedTensor, which now has one less dimension.
  
language/conpono/reconstruct/preprocess.py