def testToFromComponents(self, shape, fields, field_specs): struct = StructuredTensor.from_fields(fields, shape) spec = StructuredTensorSpec(shape, field_specs) actual_components = spec._to_components(struct) self.assertLen(actual_components, 3) self.assertAllTensorsEqual(actual_components[0], fields) rt_reconstructed = spec._from_components(actual_components) self.assertAllEqual(struct, rt_reconstructed)
def testConstruction(self): spec1_fields = dict(a=T_1_2_3_4) spec1 = StructuredTensorSpec([1, 2, 3], spec1_fields) self.assertEqual(spec1._shape, (1, 2, 3)) self.assertEqual(spec1._field_specs, spec1_fields) spec2_fields = dict(a=T_1_2, b=T_1_2_8, c=R_1_N, d=R_1_N_N, s=spec1) spec2 = StructuredTensorSpec([1, 2], spec2_fields) self.assertEqual(spec2._shape, (1, 2)) self.assertEqual(spec2._field_specs, spec2_fields)
class StructuredTensorSpecTest(test_util.TensorFlowTestCase, parameterized.TestCase): # TODO(edloper): Add a subclass of TensorFlowTestCase that overrides # assertAllEqual etc to work with StructuredTensors. def assertAllEqual(self, a, b, msg=None): if not (isinstance(a, structured_tensor.StructuredTensor) or isinstance(b, structured_tensor.StructuredTensor)): return super(StructuredTensorSpecTest, self).assertAllEqual(a, b, msg) if not (isinstance(a, structured_tensor.StructuredTensor) and isinstance(b, structured_tensor.StructuredTensor)): # TODO(edloper) Add support for this once structured_factory_ops is added. raise ValueError('Not supported yet') self.assertEqual(repr(a.shape), repr(b.shape)) self.assertEqual(set(a.field_names()), set(b.field_names())) for field in a.field_names(): self.assertAllEqual(a.field_value(field), b.field_value(field)) def assertAllTensorsEqual(self, x, y): assert isinstance(x, dict) and isinstance(y, dict) self.assertEqual(set(x), set(y)) for key in x: self.assertAllEqual(x[key], y[key]) def testConstruction(self): spec1_fields = dict(a=T_1_2_3_4) spec1 = StructuredTensorSpec([1, 2, 3], spec1_fields) self.assertEqual(spec1._shape, (1, 2, 3)) self.assertEqual(spec1._field_specs, spec1_fields) spec2_fields = dict(a=T_1_2, b=T_1_2_8, c=R_1_N, d=R_1_N_N, s=spec1) spec2 = StructuredTensorSpec([1, 2], spec2_fields) self.assertEqual(spec2._shape, (1, 2)) self.assertEqual(spec2._field_specs, spec2_fields) @parameterized.parameters([ (None, {}, r"StructuredTensor's shape must have known rank\."), ([], None, r'field_specs must be a dictionary\.'), ([], { 1: tensor_spec.TensorSpec(None) }, r'field_specs must be a dictionary with string keys\.'), ([], { 'x': 0 }, r'field_specs must be a dictionary with TypeSpec values\.'), ]) def testConstructionErrors(self, shape, field_specs, error): with self.assertRaisesRegex(TypeError, error): structured_tensor.StructuredTensorSpec(shape, field_specs) def testValueType(self): spec1 = StructuredTensorSpec([1, 2, 3], dict(a=T_1_2)) self.assertEqual(spec1.value_type, StructuredTensor) @parameterized.parameters([ (StructuredTensorSpec([1, 2, 3], {}), (tensor_shape.TensorShape([1, 2, 3]), {})), (StructuredTensorSpec([], {'a': T_1_2}), (tensor_shape.TensorShape([]), { 'a': T_1_2 })), (StructuredTensorSpec([1, 2], { 'a': T_1_2, 'b': R_1_N }), (tensor_shape.TensorShape([1, 2]), { 'a': T_1_2, 'b': R_1_N })), (StructuredTensorSpec([], {'a': T_1_2}), (tensor_shape.TensorShape([]), { 'a': T_1_2 })), ]) # pyformat: disable def testSerialize(self, spec, expected): serialization = spec._serialize() # Note that we can only use assertEqual because none of our cases include # a None dimension. A TensorShape with a None dimension is never equal # to another TensorShape. self.assertEqual(serialization, expected) @parameterized.parameters([ (StructuredTensorSpec([1, 2, 3], {}), ({}, NROWS_SPEC, (PARTITION_SPEC, PARTITION_SPEC))), (StructuredTensorSpec([], {'a': T_1_2}), ({ 'a': T_1_2 }, (), ())), (StructuredTensorSpec([1, 2], { 'a': T_1_2, 'b': R_1_N }), ({ 'a': T_1_2, 'b': R_1_N }, NROWS_SPEC, (PARTITION_SPEC, ))), (StructuredTensorSpec([], {'a': T_1_2}), ({ 'a': T_1_2 }, (), ())), ]) # pyformat: disable def testComponentSpecs(self, spec, expected): self.assertEqual(spec._component_specs, expected) @parameterized.parameters([ { 'shape': [], 'fields': dict(x=[[1.0, 2.0]]), 'field_specs': dict(x=T_1_2), }, { 'shape': [2], 'fields': dict(a=ragged_factory_ops.constant_value([[1.0], [2.0, 3.0]]), b=[[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), 'field_specs': dict(a=R_1_N, b=T_2_3), }, ]) # pyformat: disable def testToFromComponents(self, shape, fields, field_specs): struct = StructuredTensor.from_fields(fields, shape) spec = StructuredTensorSpec(shape, field_specs) actual_components = spec._to_components(struct) self.assertLen(actual_components, 3) self.assertAllTensorsEqual(actual_components[0], fields) rt_reconstructed = spec._from_components(actual_components) self.assertAllEqual(struct, rt_reconstructed) def testToFromComponentsEmptyScalar(self): struct = StructuredTensor.from_fields(fields={}, shape=[]) spec = struct._type_spec components = spec._to_components(struct) rt_reconstructed = spec._from_components(components) self.assertAllEqual(struct, rt_reconstructed) self.assertEqual(components, ({}, (), ())) def testToFromComponentsEmptyTensor(self): struct = StructuredTensor.from_fields(fields={}, shape=[1, 2, 3]) spec = struct._type_spec components = spec._to_components(struct) rt_reconstructed = spec._from_components(components) self.assertAllEqual(struct, rt_reconstructed) self.assertLen(components, 3) fields, nrows, row_partitions = components self.assertEmpty(fields) self.assertAllEqual(nrows, 1) self.assertLen(row_partitions, 2) self.assertIsInstance(row_partitions[0], row_partition.RowPartition) self.assertIsInstance(row_partitions[1], row_partition.RowPartition) self.assertAllEqual(row_partitions[0].row_splits(), [0, 2]) self.assertAllEqual(row_partitions[1].row_splits(), [0, 3, 6]) @parameterized.parameters([{ 'unbatched': StructuredTensorSpec([], {}), 'batch_size': 5, 'batched': StructuredTensorSpec([5], {}), }, { 'unbatched': StructuredTensorSpec([1, 2], {}), 'batch_size': 5, 'batched': StructuredTensorSpec([5, 1, 2], {}), }, { 'unbatched': StructuredTensorSpec([], dict(a=T_3, b=R_1_N)), 'batch_size': 2, 'batched': StructuredTensorSpec([2], dict(a=T_2_3, b=R_2_1_N)), }]) # pyformat: disable def testBatchUnbatch(self, unbatched, batch_size, batched): self.assertEqual(unbatched._batch(batch_size), batched) self.assertEqual(batched._unbatch(), unbatched) @parameterized.parameters([ { 'unbatched': lambda: [ StructuredTensor.from_fields({ 'a': 1, 'b': [5, 6] }), StructuredTensor.from_fields({ 'a': 2, 'b': [7, 8] }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields(shape=[2], fields={ 'a': [1, 2], 'b': [[5, 6], [7, 8]] }), }, { 'unbatched': lambda: [ StructuredTensor.from_fields(shape=[3], fields={ 'a': [1, 2, 3], 'b': [[5, 6], [6, 7], [7, 8]] }), StructuredTensor.from_fields(shape=[3], fields={ 'a': [2, 3, 4], 'b': [[2, 2], [3, 3], [4, 4]] }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields( shape=[2, 3], fields={ 'a': [[1, 2, 3], [2, 3, 4]], 'b': [[[5, 6], [6, 7], [7, 8]], [[2, 2], [3, 3], [4, 4]]] }), }, { 'unbatched': lambda: [ StructuredTensor.from_fields( shape=[], fields={ 'a': 1, 'b': StructuredTensor.from_fields({'x': [5]}) }), StructuredTensor.from_fields( shape=[], fields={ 'a': 2, 'b': StructuredTensor.from_fields({'x': [6]}) }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields( shape=[2], fields={ 'a': [1, 2], 'b': StructuredTensor.from_fields(shape=[2], fields={'x': [[5], [6]]}) }), }, { 'unbatched': lambda: [ StructuredTensor.from_fields( shape=[], fields={ 'Ragged3d': ragged_factory_ops.constant_value([[1, 2], [3]]), 'Ragged2d': ragged_factory_ops.constant_value([1]), }), StructuredTensor.from_fields( shape=[], fields={ 'Ragged3d': ragged_factory_ops.constant_value([[1]]), 'Ragged2d': ragged_factory_ops.constant_value([2, 3]), }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields( shape=[2], fields={ 'Ragged3d': ragged_factory_ops.constant_value([[[1, 2], [3]], [[1]]]), 'Ragged2d': ragged_factory_ops.constant_value([[1], [2, 3]]), }), 'use_only_batched_spec': True, }, ]) # pyformat: disable def testBatchUnbatchValues(self, unbatched, batch_size, batched, use_only_batched_spec=False): batched = batched() # Deferred init because it creates tensors. unbatched = unbatched() # Deferred init because it creates tensors. # Test batching. if use_only_batched_spec: unbatched_spec = type_spec.type_spec_from_value(batched)._unbatch() else: unbatched_spec = type_spec.type_spec_from_value(unbatched[0]) unbatched_tensor_lists = [ unbatched_spec._to_tensor_list(st) for st in unbatched ] batched_tensor_list = [ array_ops.stack(tensors) for tensors in zip(*unbatched_tensor_lists) ] actual_batched = unbatched_spec._batch(batch_size)._from_tensor_list( batched_tensor_list) self.assertTrue( unbatched_spec._batch(batch_size).is_compatible_with( actual_batched)) self.assertAllEqual(actual_batched, batched) # Test unbatching batched_spec = type_spec.type_spec_from_value(batched) batched_tensor_list = batched_spec._to_batched_tensor_list(batched) unbatched_tensor_lists = zip( *[array_ops.unstack(tensor) for tensor in batched_tensor_list]) actual_unbatched = [ batched_spec._unbatch()._from_tensor_list(tensor_list) for tensor_list in unbatched_tensor_lists ] self.assertLen(actual_unbatched, len(unbatched)) for st in actual_unbatched: self.assertTrue(batched_spec._unbatch().is_compatible_with(st)) for (actual, expected) in zip(actual_unbatched, unbatched): self.assertAllEqual(actual, expected)
def testValueType(self): spec1 = StructuredTensorSpec([1, 2, 3], dict(a=T_1_2)) self.assertEqual(spec1.value_type, StructuredTensor)
class StructuredTensorSpecTest(test_util.TensorFlowTestCase, parameterized.TestCase): # TODO(edloper): Add a subclass of TensorFlowTestCase that overrides # assertAllEqual etc to work with StructuredTensors. def assertAllEqual(self, a, b, msg=None): if not (isinstance(a, structured_tensor.StructuredTensor) or isinstance(b, structured_tensor.StructuredTensor)): return super(StructuredTensorSpecTest, self).assertAllEqual(a, b, msg) if not (isinstance(a, structured_tensor.StructuredTensor) and isinstance(b, structured_tensor.StructuredTensor)): # TODO(edloper) Add support for this once structured_factory_ops is added. raise ValueError('Not supported yet') self.assertEqual(repr(a.shape), repr(b.shape)) self.assertEqual(set(a.field_names()), set(b.field_names())) for field in a.field_names(): self.assertAllEqual(a.field_value(field), b.field_value(field)) def assertAllTensorsEqual(self, list1, list2): self.assertLen(list1, len(list2)) for (t1, t2) in zip(list1, list2): self.assertAllEqual(t1, t2) def testConstruction(self): spec1_fields = dict(a=T_1_2_3_4) spec1 = StructuredTensorSpec([1, 2, 3], spec1_fields) self.assertEqual(spec1._shape, (1, 2, 3)) self.assertEqual(spec1._field_specs, spec1_fields) spec2_fields = dict(a=T_1_2, b=T_1_2_8, c=R_1_N, d=R_1_N_N, s=spec1) spec2 = StructuredTensorSpec([1, 2], spec2_fields) self.assertEqual(spec2._shape, (1, 2)) self.assertEqual(spec2._field_specs, spec2_fields) def testValueType(self): spec1 = StructuredTensorSpec([1, 2, 3], dict(a=T_1_2)) self.assertEqual(spec1.value_type, StructuredTensor) @parameterized.parameters([ (StructuredTensorSpec([1, 2, 3], {}), (tensor_shape.TensorShape([1, 2, 3]), {})), (StructuredTensorSpec([], {'a': T_1_2}), (tensor_shape.TensorShape([]), { 'a': T_1_2 })), (StructuredTensorSpec([1, 2], { 'a': T_1_2, 'b': R_1_N }), (tensor_shape.TensorShape([1, 2]), { 'a': T_1_2, 'b': R_1_N })), (StructuredTensorSpec([], {'a': T_1_2}), (tensor_shape.TensorShape([]), { 'a': T_1_2 })), ]) # pyformat: disable def testSerialize(self, spec, expected): serialization = spec._serialize() # TensorShape has an unconventional definition of equality, so we can't use # assertEqual directly here. But repr() is deterministic and lossless for # the expected values, so we can use that instead. self.assertEqual(repr(serialization), repr(expected)) @parameterized.parameters([ (StructuredTensorSpec([1, 2, 3], {}), {}), (StructuredTensorSpec([], {'a': T_1_2}), { 'a': T_1_2 }), (StructuredTensorSpec([1, 2], { 'a': T_1_2, 'b': R_1_N }), { 'a': T_1_2, 'b': R_1_N }), (StructuredTensorSpec([], {'a': T_1_2}), { 'a': T_1_2 }), ]) # pyformat: disable def testComponentSpecs(self, spec, expected): self.assertEqual(spec._component_specs, expected) @parameterized.parameters([ { 'shape': [], 'fields': dict(x=[[1.0, 2.0]]), 'field_specs': dict(x=T_1_2), }, { 'shape': [1, 2, 3], 'fields': {}, 'field_specs': {}, }, { 'shape': [2], 'fields': dict(a=ragged_factory_ops.constant_value([[1.0], [2.0, 3.0]]), b=[[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), 'field_specs': dict(a=R_1_N, b=T_2_3), }, ]) # pyformat: disable def testToFromComponents(self, shape, fields, field_specs): components = fields struct = StructuredTensor(shape, fields) spec = StructuredTensorSpec(shape, field_specs) actual_components = spec._to_components(struct) self.assertAllTensorsEqual(actual_components, components) rt_reconstructed = spec._from_components(actual_components) self.assertAllEqual(struct, rt_reconstructed) @parameterized.parameters([{ 'unbatched': StructuredTensorSpec([], {}), 'batch_size': 5, 'batched': StructuredTensorSpec([5], {}), }, { 'unbatched': StructuredTensorSpec([1, 2], {}), 'batch_size': 5, 'batched': StructuredTensorSpec([5, 1, 2], {}), }, { 'unbatched': StructuredTensorSpec([], dict(a=T_3, b=R_1_N)), 'batch_size': 2, 'batched': StructuredTensorSpec([2], dict(a=T_2_3, b=R_2_1_N)), }]) # pyformat: disable def testBatchUnbatch(self, unbatched, batch_size, batched): self.assertEqual(unbatched._batch(batch_size), batched) self.assertEqual(batched._unbatch(), unbatched) @parameterized.parameters([ { 'unbatched': lambda: [ StructuredTensor([], { 'a': 1, 'b': [5, 6] }), StructuredTensor([], { 'a': 2, 'b': [7, 8] }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor([2], { 'a': [1, 2], 'b': [[5, 6], [7, 8]] }), }, { 'unbatched': lambda: [ StructuredTensor([3], { 'a': [1, 2, 3], 'b': [[5, 6], [6, 7], [7, 8]] }), StructuredTensor([3], { 'a': [2, 3, 4], 'b': [[2, 2], [3, 3], [4, 4]] }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor( [2, 3], { 'a': [[1, 2, 3], [2, 3, 4]], 'b': [[[5, 6], [6, 7], [7, 8]], [[2, 2], [3, 3], [4, 4]]] }), }, { 'unbatched': lambda: [ StructuredTensor([], { 'a': 1, 'b': StructuredTensor([], {'x': [5]}) }), StructuredTensor([], { 'a': 2, 'b': StructuredTensor([], {'x': [6]}) }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor([2], { 'a': [1, 2], 'b': StructuredTensor([2], {'x': [[5], [6]]}) }), }, ]) # pyformat: disable def testBatchUnbatchValues(self, unbatched, batch_size, batched): batched = batched() # Deferred init because it creates tensors. unbatched = unbatched() # Deferred init because it creates tensors. # Test batching. unbatched_spec = type_spec.type_spec_from_value(unbatched[0]) unbatched_tensor_lists = [ unbatched_spec._to_tensor_list(st) for st in unbatched ] batched_tensor_list = [ array_ops.stack(tensors) for tensors in zip(*unbatched_tensor_lists) ] actual_batched = unbatched_spec._batch(batch_size)._from_tensor_list( batched_tensor_list) self.assertAllEqual(actual_batched, batched) # Test unbatching batched_spec = type_spec.type_spec_from_value(batched) batched_tensor_list = batched_spec._to_tensor_list(batched) unbatched_tensor_lists = zip( *[array_ops.unstack(tensor) for tensor in batched_tensor_list]) actual_unbatched = [ batched_spec._unbatch()._from_tensor_list(tensor_list) for tensor_list in unbatched_tensor_lists ] self.assertLen(actual_unbatched, len(unbatched)) for (actual, expected) in zip(actual_unbatched, unbatched): self.assertAllEqual(actual, expected)
class StructuredTensorSpecTest(test_util.TensorFlowTestCase, parameterized.TestCase): # TODO(edloper): Add a subclass of TensorFlowTestCase that overrides # assertAllEqual etc to work with StructuredTensors. def assertAllEqual(self, a, b, msg=None): if not (isinstance(a, structured_tensor.StructuredTensor) or isinstance(b, structured_tensor.StructuredTensor)): return super(StructuredTensorSpecTest, self).assertAllEqual(a, b, msg) if not (isinstance(a, structured_tensor.StructuredTensor) and isinstance(b, structured_tensor.StructuredTensor)): # TODO(edloper) Add support for this once structured_factory_ops is added. raise ValueError('Not supported yet') self.assertEqual(repr(a.shape), repr(b.shape)) self.assertEqual(set(a.field_names()), set(b.field_names())) for field in a.field_names(): self.assertAllEqual(a.field_value(field), b.field_value(field)) def assertAllTensorsEqual(self, x, y): assert isinstance(x, dict) and isinstance(y, dict) self.assertEqual(set(x), set(y)) for key in x: self.assertAllEqual(x[key], y[key]) def testConstruction(self): spec1_fields = dict(a=T_1_2_3_4) spec1 = StructuredTensorSpec([1, 2, 3], spec1_fields) self.assertEqual(spec1._shape, (1, 2, 3)) self.assertEqual(spec1._field_specs, spec1_fields) spec2_fields = dict(a=T_1_2, b=T_1_2_8, c=R_1_N, d=R_1_N_N, s=spec1) spec2 = StructuredTensorSpec([1, 2], spec2_fields) self.assertEqual(spec2._shape, (1, 2)) self.assertEqual(spec2._field_specs, spec2_fields) @parameterized.parameters([ (None, {}, r"StructuredTensor's shape must have known rank\."), ([], None, r'field_specs must be a dictionary\.'), ([], { 1: tensor_spec.TensorSpec(None) }, r'field_specs must be a dictionary with string keys\.'), ([], { 'x': 0 }, r'field_specs must be a dictionary with TypeSpec values\.'), ]) def testConstructionErrors(self, shape, field_specs, error): with self.assertRaisesRegex(TypeError, error): structured_tensor.StructuredTensorSpec(shape, field_specs) def testValueType(self): spec1 = StructuredTensorSpec([1, 2], dict(a=T_1_2)) self.assertEqual(spec1.value_type, StructuredTensor) @parameterized.parameters([ (StructuredTensorSpec([1, 2, 3], {}), (('_fields', {}), ('_ragged_shape', dynamic_ragged_shape.DynamicRaggedShape.Spec._from_tensor_shape( [1, 2, 3], num_row_partitions=0, dtype=dtypes.int32)))), (StructuredTensorSpec([], {'a': T_1_2}), (('_fields', { 'a': T_1_2 }), ('_ragged_shape', dynamic_ragged_shape.DynamicRaggedShape.Spec._from_tensor_shape( [], num_row_partitions=0, dtype=dtypes.int64)))), (StructuredTensorSpec([1, 2], { 'a': T_1_2, 'b': R_1_N }), (('_fields', { 'a': T_1_2, 'b': R_1_N }), ('_ragged_shape', dynamic_ragged_shape.DynamicRaggedShape.Spec._from_tensor_shape( [1, 2], num_row_partitions=1, dtype=dtypes.int64)))), ]) # pyformat: disable def testSerialize(self, spec, expected): serialization = spec._serialize() # Note that we can only use assertEqual because none of our cases include # a None dimension. A TensorShape with a None dimension is never equal # to another TensorShape. self.assertEqual(serialization, expected) @parameterized.parameters([ (StructuredTensorSpec([1, 2, 3], {}), (dynamic_ragged_shape.DynamicRaggedShape.Spec._from_tensor_shape( [1, 2, 3], num_row_partitions=0, dtype=dtypes.int32), )), (StructuredTensorSpec([], {'a': T_1_2}), ( tensor_spec.TensorSpec(shape=(1, 2), dtype=dtypes.float32, name=None), dynamic_ragged_shape.DynamicRaggedShape.Spec._from_tensor_shape( [], num_row_partitions=0, dtype=dtypes.int64), )), (StructuredTensorSpec([1, 2], { 'a': T_1_2, 'b': R_1_N }), (T_1_2, R_1_N, dynamic_ragged_shape.DynamicRaggedShape.Spec._from_tensor_shape( [1, 2], num_row_partitions=1, dtype=dtypes.int64))), ]) # pyformat: disable def testComponentSpecs(self, spec, expected): self.assertEqual(spec._component_specs, expected) @parameterized.parameters([ { 'shape': [], 'fields': dict(x=[[1.0, 2.0]]), 'field_specs': dict(x=T_1_2), }, { 'shape': [2], 'fields': dict(a=ragged_factory_ops.constant_value([[1.0], [2.0, 3.0]]), b=[[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), 'field_specs': dict(a=R_2_N, b=T_2_3), }, ]) # pyformat: disable def testToFromComponents(self, shape, fields, field_specs): struct = StructuredTensor.from_fields(fields, shape) spec = StructuredTensorSpec(shape, field_specs) actual_components = spec._to_components(struct) rt_reconstructed = spec._from_components(actual_components) self.assertAllEqual(struct, rt_reconstructed) def testToFromComponentsEmptyScalar(self): struct = StructuredTensor.from_fields(fields={}, shape=[]) spec = struct._type_spec components = spec._to_components(struct) rt_reconstructed = spec._from_components(components) self.assertAllEqual(struct, rt_reconstructed) def testToFromComponentsEmptyTensor(self): struct = StructuredTensor.from_fields(fields={}, shape=[1, 2, 3]) spec = struct._type_spec components = spec._to_components(struct) rt_reconstructed = spec._from_components(components) self.assertAllEqual(struct, rt_reconstructed) @parameterized.parameters([{ 'unbatched': StructuredTensorSpec([], {}), 'batch_size': 5, 'batched': StructuredTensorSpec([5], {}), }, { 'unbatched': StructuredTensorSpec([1, 2], {}), 'batch_size': 5, 'batched': StructuredTensorSpec([5, 1, 2], {}), }, { 'unbatched': StructuredTensorSpec([], dict(a=T_3, b=R_1_N)), 'batch_size': 2, 'batched': StructuredTensorSpec([2], dict(a=T_2_3, b=R_2_1_N)), }]) # pyformat: disable def testBatchUnbatch(self, unbatched, batch_size, batched): self.assertEqual(unbatched._batch(batch_size), batched) self.assertEqual(batched._unbatch(), unbatched) @parameterized.parameters([ { 'unbatched': lambda: [ StructuredTensor.from_fields({ 'a': 1, 'b': [5, 6] }), StructuredTensor.from_fields({ 'a': 2, 'b': [7, 8] }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields(shape=[2], fields={ 'a': [1, 2], 'b': [[5, 6], [7, 8]] }), }, { 'unbatched': lambda: [ StructuredTensor.from_fields(shape=[3], fields={ 'a': [1, 2, 3], 'b': [[5, 6], [6, 7], [7, 8]] }), StructuredTensor.from_fields(shape=[3], fields={ 'a': [2, 3, 4], 'b': [[2, 2], [3, 3], [4, 4]] }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields( shape=[2, 3], fields={ 'a': [[1, 2, 3], [2, 3, 4]], 'b': [[[5, 6], [6, 7], [7, 8]], [[2, 2], [3, 3], [4, 4]]] }), }, { 'unbatched': lambda: [ StructuredTensor.from_fields( shape=[], fields={ 'a': 1, 'b': StructuredTensor.from_fields({'x': [5]}) }), StructuredTensor.from_fields( shape=[], fields={ 'a': 2, 'b': StructuredTensor.from_fields({'x': [6]}) }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields( shape=[2], fields={ 'a': [1, 2], 'b': StructuredTensor.from_fields(shape=[2], fields={'x': [[5], [6]]}) }), }, { 'unbatched': lambda: [ StructuredTensor.from_fields( shape=[], fields={ 'Ragged3d': ragged_factory_ops.constant_value([[1, 2], [3]]), 'Ragged2d': ragged_factory_ops.constant_value([1]), }), StructuredTensor.from_fields( shape=[], fields={ 'Ragged3d': ragged_factory_ops.constant_value([[1]]), 'Ragged2d': ragged_factory_ops.constant_value([2, 3]), }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields( shape=[2], fields={ 'Ragged3d': ragged_factory_ops.constant_value([[[1, 2], [3]], [[1]]]), 'Ragged2d': ragged_factory_ops.constant_value([[1], [2, 3]]), }), 'use_only_batched_spec': True, }, ]) # pyformat: disable def testBatchUnbatchValues(self, unbatched, batch_size, batched, use_only_batched_spec=False): batched = batched() # Deferred init because it creates tensors. unbatched = unbatched() # Deferred init because it creates tensors. def unbatch_gen(): for i in unbatched: yield i ds = dataset_ops.Dataset.from_tensors(batched) ds2 = ds.unbatch() if context.executing_eagerly(): v = list(ds2.batch(2)) self.assertAllEqual(v[0], batched) if not use_only_batched_spec: unbatched_spec = type_spec.type_spec_from_value(unbatched[0]) dsu = dataset_ops.Dataset.from_generator( unbatch_gen, output_signature=unbatched_spec) dsu2 = dsu.batch(2) if context.executing_eagerly(): v = list(dsu2) self.assertAllEqual(v[0], batched) def _lambda_for_fields(self): return lambda: { 'a': np.ones([1, 2, 3, 1]), 'b': np.ones([1, 2, 3, 1, 5]), 'c': ragged_factory_ops.constant(np.zeros([1, 2, 3, 1], dtype=np.uint8), dtype=dtypes.uint8), 'd': ragged_factory_ops.constant(np.zeros([1, 2, 3, 1, 3]).tolist(), ragged_rank=1), 'e': ragged_factory_ops.constant(np.zeros([1, 2, 3, 1, 2, 2]).tolist(), ragged_rank=2), 'f': ragged_factory_ops.constant(np.zeros([1, 2, 3, 1, 3]), dtype=dtypes.float32), 'g': StructuredTensor.from_pyval([[ [ # pylint: disable=g-complex-comprehension [{ 'x': j, 'y': k }] for k in range(3) ] for j in range(2) ]]), 'h': StructuredTensor.from_pyval([[ [ # pylint: disable=g-complex-comprehension [[{ 'x': j, 'y': k, 'z': z } for z in range(j)]] for k in range(3) ] for j in range(2) ]]), }
class StructuredTensorSpecTest(test_util.TensorFlowTestCase, parameterized.TestCase): # TODO(edloper): Add a subclass of TensorFlowTestCase that overrides # assertAllEqual etc to work with StructuredTensors. def assertAllEqual(self, a, b, msg=None): if not (isinstance(a, structured_tensor.StructuredTensor) or isinstance(b, structured_tensor.StructuredTensor)): return super(StructuredTensorSpecTest, self).assertAllEqual(a, b, msg) if not (isinstance(a, structured_tensor.StructuredTensor) and isinstance(b, structured_tensor.StructuredTensor)): # TODO(edloper) Add support for this once structured_factory_ops is added. raise ValueError('Not supported yet') self.assertEqual(repr(a.shape), repr(b.shape)) self.assertEqual(set(a.field_names()), set(b.field_names())) for field in a.field_names(): self.assertAllEqual(a.field_value(field), b.field_value(field)) def assertAllTensorsEqual(self, x, y): assert isinstance(x, dict) and isinstance(y, dict) self.assertEqual(set(x), set(y)) for key in x: self.assertAllEqual(x[key], y[key]) def testConstruction(self): spec1_fields = dict(a=T_1_2_3_4) spec1 = StructuredTensorSpec([1, 2, 3], spec1_fields) self.assertEqual(spec1._shape, (1, 2, 3)) self.assertEqual(spec1._field_specs, spec1_fields) spec2_fields = dict(a=T_1_2, b=T_1_2_8, c=R_1_N, d=R_1_N_N, s=spec1) spec2 = StructuredTensorSpec([1, 2], spec2_fields) self.assertEqual(spec2._shape, (1, 2)) self.assertEqual(spec2._field_specs, spec2_fields) @parameterized.parameters([ (None, {}, r"StructuredTensor's shape must have known rank\."), ([], None, r'field_specs must be a dictionary\.'), ([], { 1: tensor_spec.TensorSpec(None) }, r'field_specs must be a dictionary with string keys\.'), ([], { 'x': 0 }, r'field_specs must be a dictionary with TypeSpec values\.'), ]) def testConstructionErrors(self, shape, field_specs, error): with self.assertRaisesRegex(TypeError, error): structured_tensor.StructuredTensorSpec(shape, field_specs) def testValueType(self): spec1 = StructuredTensorSpec([1, 2], dict(a=T_1_2)) self.assertEqual(spec1.value_type, StructuredTensor) @parameterized.parameters([ (StructuredTensorSpec([1, 2, 3], {}), (tensor_shape.TensorShape([1, 2, 3]), {})), (StructuredTensorSpec([], {'a': T_1_2}), (tensor_shape.TensorShape([]), { 'a': T_1_2 })), (StructuredTensorSpec([1, 2], { 'a': T_1_2, 'b': R_1_N }), (tensor_shape.TensorShape([1, 2]), { 'a': T_1_2, 'b': R_1_N })), (StructuredTensorSpec([], {'a': T_1_2}), (tensor_shape.TensorShape([]), { 'a': T_1_2 })), ]) # pyformat: disable def testSerialize(self, spec, expected): serialization = spec._serialize() # Note that we can only use assertEqual because none of our cases include # a None dimension. A TensorShape with a None dimension is never equal # to another TensorShape. self.assertEqual(serialization, expected) @parameterized.parameters([ (StructuredTensorSpec([1, 2, 3], {}), ({}, NROWS_SPEC, (PARTITION_SPEC, PARTITION_SPEC))), (StructuredTensorSpec([], {'a': T_1_2}), ({ 'a': T_1_2 }, (), ())), (StructuredTensorSpec([1, 2], { 'a': T_1_2, 'b': R_1_N }), ({ 'a': T_1_2, 'b': R_1_N }, NROWS_SPEC, (PARTITION_SPEC, ))), (StructuredTensorSpec([], {'a': T_1_2}), ({ 'a': T_1_2 }, (), ())), ]) # pyformat: disable def testComponentSpecs(self, spec, expected): self.assertEqual(spec._component_specs, expected) @parameterized.parameters([ { 'shape': [], 'fields': dict(x=[[1.0, 2.0]]), 'field_specs': dict(x=T_1_2), }, { 'shape': [2], 'fields': dict(a=ragged_factory_ops.constant_value([[1.0], [2.0, 3.0]]), b=[[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), 'field_specs': dict(a=R_1_N, b=T_2_3), }, ]) # pyformat: disable def testToFromComponents(self, shape, fields, field_specs): struct = StructuredTensor.from_fields(fields, shape) spec = StructuredTensorSpec(shape, field_specs) actual_components = spec._to_components(struct) self.assertLen(actual_components, 3) self.assertAllTensorsEqual(actual_components[0], fields) rt_reconstructed = spec._from_components(actual_components) self.assertAllEqual(struct, rt_reconstructed) def testToFromComponentsEmptyScalar(self): struct = StructuredTensor.from_fields(fields={}, shape=[]) spec = struct._type_spec components = spec._to_components(struct) rt_reconstructed = spec._from_components(components) self.assertAllEqual(struct, rt_reconstructed) self.assertEqual(components, ({}, (), ())) def testToFromComponentsEmptyTensor(self): struct = StructuredTensor.from_fields(fields={}, shape=[1, 2, 3]) spec = struct._type_spec components = spec._to_components(struct) rt_reconstructed = spec._from_components(components) self.assertAllEqual(struct, rt_reconstructed) self.assertLen(components, 3) fields, nrows, row_partitions = components self.assertEmpty(fields) self.assertAllEqual(nrows, 1) self.assertLen(row_partitions, 2) self.assertIsInstance(row_partitions[0], row_partition.RowPartition) self.assertIsInstance(row_partitions[1], row_partition.RowPartition) self.assertAllEqual(row_partitions[0].row_splits(), [0, 2]) self.assertAllEqual(row_partitions[1].row_splits(), [0, 3, 6]) @parameterized.parameters([{ 'unbatched': StructuredTensorSpec([], {}), 'batch_size': 5, 'batched': StructuredTensorSpec([5], {}), }, { 'unbatched': StructuredTensorSpec([1, 2], {}), 'batch_size': 5, 'batched': StructuredTensorSpec([5, 1, 2], {}), }, { 'unbatched': StructuredTensorSpec([], dict(a=T_3, b=R_1_N)), 'batch_size': 2, 'batched': StructuredTensorSpec([2], dict(a=T_2_3, b=R_2_1_N)), }]) # pyformat: disable def testBatchUnbatch(self, unbatched, batch_size, batched): self.assertEqual(unbatched._batch(batch_size), batched) self.assertEqual(batched._unbatch(), unbatched) @parameterized.parameters([ { 'unbatched': lambda: [ StructuredTensor.from_fields({ 'a': 1, 'b': [5, 6] }), StructuredTensor.from_fields({ 'a': 2, 'b': [7, 8] }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields(shape=[2], fields={ 'a': [1, 2], 'b': [[5, 6], [7, 8]] }), }, { 'unbatched': lambda: [ StructuredTensor.from_fields(shape=[3], fields={ 'a': [1, 2, 3], 'b': [[5, 6], [6, 7], [7, 8]] }), StructuredTensor.from_fields(shape=[3], fields={ 'a': [2, 3, 4], 'b': [[2, 2], [3, 3], [4, 4]] }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields( shape=[2, 3], fields={ 'a': [[1, 2, 3], [2, 3, 4]], 'b': [[[5, 6], [6, 7], [7, 8]], [[2, 2], [3, 3], [4, 4]]] }), }, { 'unbatched': lambda: [ StructuredTensor.from_fields( shape=[], fields={ 'a': 1, 'b': StructuredTensor.from_fields({'x': [5]}) }), StructuredTensor.from_fields( shape=[], fields={ 'a': 2, 'b': StructuredTensor.from_fields({'x': [6]}) }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields( shape=[2], fields={ 'a': [1, 2], 'b': StructuredTensor.from_fields(shape=[2], fields={'x': [[5], [6]]}) }), }, { 'unbatched': lambda: [ StructuredTensor.from_fields( shape=[], fields={ 'Ragged3d': ragged_factory_ops.constant_value([[1, 2], [3]]), 'Ragged2d': ragged_factory_ops.constant_value([1]), }), StructuredTensor.from_fields( shape=[], fields={ 'Ragged3d': ragged_factory_ops.constant_value([[1]]), 'Ragged2d': ragged_factory_ops.constant_value([2, 3]), }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields( shape=[2], fields={ 'Ragged3d': ragged_factory_ops.constant_value([[[1, 2], [3]], [[1]]]), 'Ragged2d': ragged_factory_ops.constant_value([[1], [2, 3]]), }), 'use_only_batched_spec': True, }, ]) # pyformat: disable def testBatchUnbatchValues(self, unbatched, batch_size, batched, use_only_batched_spec=False): batched = batched() # Deferred init because it creates tensors. unbatched = unbatched() # Deferred init because it creates tensors. # Test batching. if use_only_batched_spec: unbatched_spec = type_spec.type_spec_from_value(batched)._unbatch() else: unbatched_spec = type_spec.type_spec_from_value(unbatched[0]) unbatched_tensor_lists = [ unbatched_spec._to_tensor_list(st) for st in unbatched ] batched_tensor_list = [ array_ops.stack(tensors) for tensors in zip(*unbatched_tensor_lists) ] actual_batched = unbatched_spec._batch(batch_size)._from_tensor_list( batched_tensor_list) self.assertTrue( unbatched_spec._batch(batch_size).is_compatible_with( actual_batched)) self.assertAllEqual(actual_batched, batched) # Test unbatching batched_spec = type_spec.type_spec_from_value(batched) batched_tensor_list = batched_spec._to_batched_tensor_list(batched) unbatched_tensor_lists = zip( *[array_ops.unstack(tensor) for tensor in batched_tensor_list]) actual_unbatched = [ batched_spec._unbatch()._from_tensor_list(tensor_list) for tensor_list in unbatched_tensor_lists ] self.assertLen(actual_unbatched, len(unbatched)) for st in actual_unbatched: self.assertTrue(batched_spec._unbatch().is_compatible_with(st)) for (actual, expected) in zip(actual_unbatched, unbatched): self.assertAllEqual(actual, expected) def _lambda_for_fields(self): return lambda: { 'a': np.ones([1, 2, 3, 1]), 'b': np.ones([1, 2, 3, 1, 5]), 'c': ragged_factory_ops.constant(np.zeros([1, 2, 3, 1], dtype=np.uint8), dtype=dtypes.uint8), 'd': ragged_factory_ops.constant(np.zeros([1, 2, 3, 1, 3]).tolist(), ragged_rank=1), 'e': ragged_factory_ops.constant(np.zeros([1, 2, 3, 1, 2, 2]).tolist(), ragged_rank=2), 'f': ragged_factory_ops.constant(np.zeros([1, 2, 3, 1, 3]), dtype=dtypes.float32), 'g': StructuredTensor.from_pyval([[ [ # pylint: disable=g-complex-comprehension [{ 'x': j, 'y': k }] for k in range(3) ] for j in range(2) ]]), 'h': StructuredTensor.from_pyval([[ [ # pylint: disable=g-complex-comprehension [[{ 'x': j, 'y': k, 'z': z } for z in range(j)]] for k in range(3) ] for j in range(2) ]]), } def testFlatTensorSpecs(self): # Note that the batchable tensor list encoding for a StructuredTensor # contains a separate tensor for each leaf field. # In this example, _flat_tensor_specs in class StructuredTensorSpec is # called three times and it returns results with length 2, 3 and 11 # for "g", "h" and `struct` respectively. fields = self._lambda_for_fields() rank = 4 if callable(fields): fields = fields( ) # deferred construction: fields may include tensors. struct = StructuredTensor.from_fields_and_rank(fields, rank) spec = type_spec.type_spec_from_value(struct) flat_specs = spec._flat_tensor_specs self.assertEqual( flat_specs, [ # a , b tensor_spec.TensorSpec( shape=(1, 2, 3, 1), dtype=dtypes.float64, name=None), tensor_spec.TensorSpec( shape=(1, 2, 3, 1, 5), dtype=dtypes.float64, name=None), # c, d, e, f tensor_spec.TensorSpec( shape=None, dtype=dtypes.variant, name=None), tensor_spec.TensorSpec( shape=None, dtype=dtypes.variant, name=None), tensor_spec.TensorSpec( shape=None, dtype=dtypes.variant, name=None), tensor_spec.TensorSpec( shape=None, dtype=dtypes.variant, name=None), # g tensor_spec.TensorSpec( shape=None, dtype=dtypes.variant, name=None), tensor_spec.TensorSpec( shape=None, dtype=dtypes.variant, name=None), # h tensor_spec.TensorSpec( shape=None, dtype=dtypes.variant, name=None), tensor_spec.TensorSpec( shape=None, dtype=dtypes.variant, name=None), tensor_spec.TensorSpec( shape=None, dtype=dtypes.variant, name=None) ]) def testFulTypesForFlatTensors(self): # Note that the batchable tensor list encoding for a StructuredTensor # contains a separate tensor for each leaf field. # In this example, _flat_tensor_specs in class StructuredTensorSpec is # called three times and it returns results with length 2, 3 and 11 # for "g", "h" and `struct` respectively. fields = self._lambda_for_fields() rank = 4 if callable(fields): fields = fields( ) # deferred construction: fields may include tensors. struct = StructuredTensor.from_fields_and_rank(fields, rank) spec = type_spec.type_spec_from_value(struct) flat_specs = spec._flat_tensor_specs fulltype = fulltypes_for_flat_tensors(spec) expected_ft_list = [ # a, b full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_UNSET), full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_UNSET), # c, d, e, f full_type_pb2.FullTypeDef( type_id=full_type_pb2.TFT_RAGGED, args=[ full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_UINT8) ]), full_type_pb2.FullTypeDef( type_id=full_type_pb2.TFT_RAGGED, args=[ full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_FLOAT) ]), full_type_pb2.FullTypeDef( type_id=full_type_pb2.TFT_RAGGED, args=[ full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_FLOAT) ]), full_type_pb2.FullTypeDef( type_id=full_type_pb2.TFT_RAGGED, args=[ full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_FLOAT) ]), # g full_type_pb2.FullTypeDef( type_id=full_type_pb2.TFT_RAGGED, args=[ full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_INT32) ]), full_type_pb2.FullTypeDef( type_id=full_type_pb2.TFT_RAGGED, args=[ full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_INT32) ]), # h full_type_pb2.FullTypeDef( type_id=full_type_pb2.TFT_RAGGED, args=[ full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_INT32) ]), full_type_pb2.FullTypeDef( type_id=full_type_pb2.TFT_RAGGED, args=[ full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_INT32) ]), full_type_pb2.FullTypeDef( type_id=full_type_pb2.TFT_RAGGED, args=[ full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_INT32) ]), ] self.assertEqual(len(expected_ft_list), len(flat_specs)) self.assertEqual(fulltype, expected_ft_list)
class StructuredTensorSpecTest(test_util.TensorFlowTestCase, parameterized.TestCase): # TODO(edloper): Add a subclass of TensorFlowTestCase that overrides # assertAllEqual etc to work with StructuredTensors. def assertAllEqual(self, a, b, msg=None): if not (isinstance(a, structured_tensor.StructuredTensor) or isinstance(b, structured_tensor.StructuredTensor)): return super(StructuredTensorSpecTest, self).assertAllEqual(a, b, msg) if not (isinstance(a, structured_tensor.StructuredTensor) and isinstance(b, structured_tensor.StructuredTensor)): # TODO(edloper) Add support for this once structured_factory_ops is added. raise ValueError('Not supported yet') self.assertEqual(repr(a.shape), repr(b.shape)) self.assertEqual(set(a.field_names()), set(b.field_names())) for field in a.field_names(): self.assertAllEqual(a.field_value(field), b.field_value(field)) def assertAllTensorsEqual(self, list1, list2): self.assertLen(list1, len(list2)) for (t1, t2) in zip(list1, list2): self.assertAllEqual(t1, t2) def testConstruction(self): spec1_fields = dict(a=T_1_2_3_4) spec1 = StructuredTensorSpec([1, 2, 3], spec1_fields) self.assertEqual(spec1._shape, (1, 2, 3)) self.assertEqual(spec1._field_specs, spec1_fields) spec2_fields = dict(a=T_1_2, b=T_1_2_8, c=R_1_N, d=R_1_N_N, s=spec1) spec2 = StructuredTensorSpec([1, 2], spec2_fields) self.assertEqual(spec2._shape, (1, 2)) self.assertEqual(spec2._field_specs, spec2_fields) @parameterized.parameters([ (None, {}, r"StructuredTensor's shape must have known rank\."), ([], None, r'field_specs must be a dictionary\.'), ([], { 1: tensor_spec.TensorSpec(None) }, r'field_specs must be a dictionary with string keys\.'), ([], { 'x': 0 }, r'field_specs must be a dictionary with TypeSpec values\.'), ]) def testConstructionErrors(self, shape, field_specs, error): with self.assertRaisesRegex(TypeError, error): structured_tensor.StructuredTensorSpec(shape, field_specs) def testValueType(self): spec1 = StructuredTensorSpec([1, 2, 3], dict(a=T_1_2)) self.assertEqual(spec1.value_type, StructuredTensor) @parameterized.parameters([ (StructuredTensorSpec([1, 2, 3], {}), (tensor_shape.TensorShape([1, 2, 3]), {})), (StructuredTensorSpec([], {'a': T_1_2}), (tensor_shape.TensorShape([]), { 'a': T_1_2 })), (StructuredTensorSpec([1, 2], { 'a': T_1_2, 'b': R_1_N }), (tensor_shape.TensorShape([1, 2]), { 'a': T_1_2, 'b': R_1_N })), (StructuredTensorSpec([], {'a': T_1_2}), (tensor_shape.TensorShape([]), { 'a': T_1_2 })), ]) # pyformat: disable def testSerialize(self, spec, expected): serialization = spec._serialize() # Note that we can only use assertEqual because none of our cases include # a None dimension. A TensorShape with a None dimension is never equal # to another TensorShape. self.assertEqual(serialization, expected) @parameterized.parameters([ (StructuredTensorSpec([1, 2, 3], {}), {}), (StructuredTensorSpec([], {'a': T_1_2}), { 'a': T_1_2 }), (StructuredTensorSpec([1, 2], { 'a': T_1_2, 'b': R_1_N }), { 'a': T_1_2, 'b': R_1_N }), (StructuredTensorSpec([], {'a': T_1_2}), { 'a': T_1_2 }), ]) # pyformat: disable def testComponentSpecs(self, spec, expected): self.assertEqual(spec._component_specs, expected) @parameterized.parameters([ { 'shape': [], 'fields': dict(x=[[1.0, 2.0]]), 'field_specs': dict(x=T_1_2), }, # TODO(edloper): Enable this test once we update StructuredTensorSpec # to contain the shared row partitions. #{ # 'shape': [1, 2, 3], # 'fields': {}, # 'field_specs': {}, #}, { 'shape': [2], 'fields': dict(a=ragged_factory_ops.constant_value([[1.0], [2.0, 3.0]]), b=[[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), 'field_specs': dict(a=R_1_N, b=T_2_3), }, ]) # pyformat: disable def testToFromComponents(self, shape, fields, field_specs): components = fields struct = StructuredTensor.from_fields(fields, shape) spec = StructuredTensorSpec(shape, field_specs) actual_components = spec._to_components(struct) self.assertAllTensorsEqual(actual_components, components) rt_reconstructed = spec._from_components(actual_components) self.assertAllEqual(struct, rt_reconstructed) @parameterized.parameters([{ 'unbatched': StructuredTensorSpec([], {}), 'batch_size': 5, 'batched': StructuredTensorSpec([5], {}), }, { 'unbatched': StructuredTensorSpec([1, 2], {}), 'batch_size': 5, 'batched': StructuredTensorSpec([5, 1, 2], {}), }, { 'unbatched': StructuredTensorSpec([], dict(a=T_3, b=R_1_N)), 'batch_size': 2, 'batched': StructuredTensorSpec([2], dict(a=T_2_3, b=R_2_1_N)), }]) # pyformat: disable def testBatchUnbatch(self, unbatched, batch_size, batched): self.assertEqual(unbatched._batch(batch_size), batched) self.assertEqual(batched._unbatch(), unbatched) @parameterized.parameters([ { 'unbatched': lambda: [ StructuredTensor.from_fields({ 'a': 1, 'b': [5, 6] }), StructuredTensor.from_fields({ 'a': 2, 'b': [7, 8] }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields(shape=[2], fields={ 'a': [1, 2], 'b': [[5, 6], [7, 8]] }), }, { 'unbatched': lambda: [ StructuredTensor.from_fields(shape=[3], fields={ 'a': [1, 2, 3], 'b': [[5, 6], [6, 7], [7, 8]] }), StructuredTensor.from_fields(shape=[3], fields={ 'a': [2, 3, 4], 'b': [[2, 2], [3, 3], [4, 4]] }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields( shape=[2, 3], fields={ 'a': [[1, 2, 3], [2, 3, 4]], 'b': [[[5, 6], [6, 7], [7, 8]], [[2, 2], [3, 3], [4, 4]]] }), }, { 'unbatched': lambda: [ StructuredTensor.from_fields( shape=[], fields={ 'a': 1, 'b': StructuredTensor.from_fields({'x': [5]}) }), StructuredTensor.from_fields( shape=[], fields={ 'a': 2, 'b': StructuredTensor.from_fields({'x': [6]}) }) ], 'batch_size': 2, 'batched': lambda: StructuredTensor.from_fields( shape=[2], fields={ 'a': [1, 2], 'b': StructuredTensor.from_fields(shape=[2], fields={'x': [[5], [6]]}) }), }, ]) # pyformat: disable def testBatchUnbatchValues(self, unbatched, batch_size, batched): batched = batched() # Deferred init because it creates tensors. unbatched = unbatched() # Deferred init because it creates tensors. # Test batching. unbatched_spec = type_spec.type_spec_from_value(unbatched[0]) unbatched_tensor_lists = [ unbatched_spec._to_tensor_list(st) for st in unbatched ] batched_tensor_list = [ array_ops.stack(tensors) for tensors in zip(*unbatched_tensor_lists) ] actual_batched = unbatched_spec._batch(batch_size)._from_tensor_list( batched_tensor_list) self.assertAllEqual(actual_batched, batched) # Test unbatching batched_spec = type_spec.type_spec_from_value(batched) batched_tensor_list = batched_spec._to_tensor_list(batched) unbatched_tensor_lists = zip( *[array_ops.unstack(tensor) for tensor in batched_tensor_list]) actual_unbatched = [ batched_spec._unbatch()._from_tensor_list(tensor_list) for tensor_list in unbatched_tensor_lists ] self.assertLen(actual_unbatched, len(unbatched)) for (actual, expected) in zip(actual_unbatched, unbatched): self.assertAllEqual(actual, expected)