def test_add_ragged_fields(self):
# Nested `Sequence(Sequence(tf.int64))`
example_data = [
[1, 2, 3],
[],
[4, 5],
]
tensor_info = feature_lib.TensorInfo(shape=(
None,
None,
),
dtype=tf.int64,
sequence_rank=2)
out = example_serializer._add_ragged_fields(example_data, tensor_info)
self.assertRaggedFieldEqual(
out, {
'ragged_flat_values': (
np.array([1, 2, 3, 4, 5]),
feature_lib.TensorInfo(shape=(None, ), dtype=tf.int64),
),
'ragged_row_lengths_0': (
[3, 0, 2],
feature_lib.TensorInfo(shape=(None, ), dtype=tf.int64),
),
})
def test_add_ragged_fields_single_level_sequence(self):
# Single level sequence
example_data = [
[1, 2],
[2, 3],
[4, 5],
]
tensor_info = feature_lib.TensorInfo(shape=(
None,
2,
),
dtype=tf.int64,
sequence_rank=1)
out = example_serializer._add_ragged_fields(example_data, tensor_info)
self.assertAllEqual(out[0], [
[1, 2],
[2, 3],
[4, 5],
])
self.assertEqual(out[1], tensor_info)
def test_add_ragged_fields_all_empty(self):
# Empty list
example_data = []
tensor_info = feature_lib.TensorInfo(shape=(
None,
None,
),
dtype=tf.int64,
sequence_rank=2)
out = example_serializer._add_ragged_fields(example_data, tensor_info)
self.assertRaggedFieldEqual(
out, {
'ragged_flat_values': (
np.zeros(shape=(0, ), dtype=np.int64),
feature_lib.TensorInfo(shape=(None, ), dtype=tf.int64),
),
'ragged_row_lengths_0': (
np.zeros(shape=(0, ), dtype=np.int64),
feature_lib.TensorInfo(shape=(None, ), dtype=tf.int64),
),
})
