def _make_composite_tensor_wrapper_if_composite(
    x: common_types.ConsistentTensorType
) -> Callable[[tf.Tensor], common_types.ConsistentTensorType]:
  """Produces a function to wrap values in the composite structure of x."""
  if isinstance(x, tf.SparseTensor):
    return lambda values: tf.SparseTensor(x.indices, values, x.dense_shape)
  elif isinstance(x, tf.RaggedTensor):

    def from_nested_row_splits(values):
      return tf.RaggedTensor.from_nested_row_splits(
          values, x.nested_row_splits, validate=False)

    return from_nested_row_splits
  else:
    return lambda values: values
Produces a function to wrap values in the composite structure of x.
tensorflow_transform/mappers.py