How to write Ragged tensors to TFRecords?

Question

I am trying to write a dataset made of ragged tensors in TFRecords, following this suggestion and the standard steps by TensorFlow, but I have issues when mapping the serialize function to the features dataset. Using a simplified case, this could be my data:

cells = tf.random_uniform(shape=([10, 3, 1]))
cells_s = tf.random_uniform(shape=([10, 4, 1]))

cells_k = tf.concat([tf.RaggedTensor.from_tensor(cells), tf.RaggedTensor.from_tensor(cells_s)], axis=0) #my data for the dataset

To serialize it, I use these functions:

def _bytes_feature(value):
  """Returns a bytes_list from a string / byte."""
  if isinstance(value, type(tf.constant(0))):
    value = value.numpy() # BytesList won't unpack a string from an EagerTensor.
  return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def serialize_example_r(feature0):
    feature = {
        'cells_val': _bytes_feature(tf.io.serialize_tensor(feature0.flat_values)),
        'cells_rs': _bytes_feature(tf.io.serialize_tensor(feature0.row_splits))
    }
    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
    return example_proto.SerializeToString()

Then I use from_tensor_slices to create the dataset and map the serialize_example_r function to the dataset to obtain the serialized features dataset:

features_dataset = tf.data.Dataset.from_tensor_slices(cells_k)

def tf_serialize_example(f0):
  tf_string = tf.py_function(
    serialize_example_r,
    (f0),  # Pass these args to the above function.
    tf.string)      # The return type is `tf.string`.
  return tf.reshape(tf_string, ()) # The result is a scalar.

serialized_features_dataset = features_dataset.map(tf_serialize_example)

But I receive this error:

ValueError: ragged_rank must be non-negative; got 0.

What am I doing wrong? What is the right process to create a TFRecords file from Ragged tensors? I am trying to put pieces together but I don't seem to find a solution. Thanks a lot in advance!

How to write Ragged tensors to TFRecords?

0 Answers