Ejemplo n.º 1
0
 def setUp(self):
     super().setUp()
     global_ids = test_util.create_device_ids_array((2, 2))
     local_device_ids = np.ravel(global_ids).tolist()
     mesh_dict = {
         'CPU':
         dtensor.Mesh(['X', 'Y'], global_ids, local_device_ids,
                      test_util.create_device_list((2, 2), 'CPU'))
     }
     self.mesh = self.configTestMesh(mesh_dict)
Ejemplo n.º 2
0
 def setUp(self):
   super().setUp()
   backend.enable_tf_random_generator()
   tf_utils.set_random_seed(1337)
   global_ids = test_util.create_device_ids_array((2, 2))
   local_device_ids = np.ravel(global_ids).tolist()
   mesh_dict = {
       'CPU':
           dtensor.Mesh(['X', 'Y'], global_ids,
                        local_device_ids,
                        test_util.create_device_list((2, 2), 'CPU'))
   }
   self.mesh = self.configTestMesh(mesh_dict)
Ejemplo n.º 3
0
 def setUp(self):
     super().setUp()
     global_ids = test_util.create_device_ids_array((2, 2))
     local_device_ids = np.ravel(global_ids).tolist()
     mesh_dict = {
         "CPU": dtensor.Mesh(
             ["X", "Y"],
             global_ids,
             local_device_ids,
             test_util.create_device_list((2, 2), "CPU"),
         )
     }
     self.mesh = self.configTestMesh(mesh_dict)