Tests freezing the queue.
(self)
| 99 | i.set_number_of_shards(4) |
| 100 | |
| 101 | def testFreezing(self): |
| 102 | """Tests freezing the queue.""" |
| 103 | i = tpu_feed.InfeedQueue(number_of_tuple_elements=2) |
| 104 | t1 = constant_op.constant(1, dtypes.int32, shape=[2]) |
| 105 | t2 = constant_op.constant(2.0, dtypes.float32, shape=[2, 4]) |
| 106 | i.set_configuration_from_sharded_input_tensors([[t2, t1], [t2, t1]]) |
| 107 | self.assertEqual(i.number_of_shards, 2) |
| 108 | self.assertEqual(i.tuple_shapes, [[4, 4], [4]]) |
| 109 | self.assertEqual(i.tuple_types, [dtypes.float32, dtypes.int32]) |
| 110 | self.assertEqual(i.shard_dimensions, [0, 0]) |
| 111 | i.freeze() |
| 112 | i.set_number_of_shards(2) |
| 113 | i.set_tuple_shapes([[4, 4], [4]]) |
| 114 | i.set_tuple_types([dtypes.float32, dtypes.int32]) |
| 115 | i.set_shard_dimensions([0, 0]) |
| 116 | with self.assertRaises(ValueError): |
| 117 | i.set_number_of_shards(1) |
| 118 | with self.assertRaises(ValueError): |
| 119 | i.set_tuple_shapes([[8, 8], [8]]) |
| 120 | with self.assertRaises(ValueError): |
| 121 | i.set_tuple_types([dtypes.int32, dtypes.float32]) |
| 122 | with self.assertRaises(ValueError): |
| 123 | i.set_shard_dimensions([1, 0]) |
| 124 | self.assertEqual(i.number_of_shards, 2) |
| 125 | self.assertEqual(i.tuple_shapes, [[4, 4], [4]]) |
| 126 | self.assertEqual(i.tuple_types, [dtypes.float32, dtypes.int32]) |
| 127 | self.assertEqual(i.shard_dimensions, [0, 0]) |
| 128 | |
| 129 | if __name__ == '__main__': |
| 130 | test.main() |
nothing calls this directly
no test coverage detected