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), ), })