def _ragged_feature_spec_from_batched_tensor(
    name: str,
    tensor: tf.RaggedTensor) -> Union[tf.io.VarLenFeature, tf.io.RaggedFeature]:
  """Infer `tf.io.RaggedFeature` from a batched `tf.RaggedTensor`."""
  partitions = []
  row_lengths_partition_idx = 1
  # Ignore batch dimension.
  for dim in tensor.values.shape[1:]:
    if dim or dim == 0:
      partitions.append(
          tf.io.RaggedFeature.UniformRowLength(  # pytype: disable=attribute-error
              length=dim))
    else:
      partitions.append(
          tf.io.RaggedFeature.RowLengths(  # pytype: disable=attribute-error
              key='{}$row_lengths_{}'.format(name, row_lengths_partition_idx)))
      row_lengths_partition_idx += 1

  return tf.io.RaggedFeature(
      dtype=tensor.dtype,
      value_key='{}$ragged_values'.format(name),
      partitions=partitions,
      row_splits_dtype=tensor.row_splits.dtype)
Infer `tf.io.RaggedFeature` from a batched `tf.RaggedTensor`.
tensorflow_transform/schema_inference.py