• Home
  • History
  • Annotate
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 # Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # =============================================================================
15 
16 """Tests for topology.py."""
17 
18 from __future__ import absolute_import
19 from __future__ import division
20 from __future__ import print_function
21 
22 from tensorflow.python.platform import test
23 from tensorflow.python.tpu import topology
24 
25 
26 class TopologyTest(test.TestCase):
27 
28   def testSerialization(self):
29     """Tests if the class is able to generate serialized strings."""
30     original_topology = topology.Topology(
31         mesh_shape=[1, 1, 2],
32         device_coordinates=[[[0, 0, 0], [0, 0, 1]]],
33     )
34     serialized_str = original_topology.serialized()
35     new_topology = topology.Topology(serialized=serialized_str)
36 
37     # Make sure the topology recovered from serialized str is same as the
38     # original topology.
39     self.assertAllEqual(
40         original_topology.mesh_shape, new_topology.mesh_shape)
41     self.assertAllEqual(
42         original_topology.device_coordinates, new_topology.device_coordinates)
43 
44 if __name__ == "__main__":
45   test.main()
46