def setUp(self): super(OptimizersTest, self).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)
def setUp(self): super(UtilsTest, self).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) self.layout = dtensor.Layout.replicated(self.mesh, rank=1)
def setUp(self): super(LayersTest, self).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)
def setUp(self): super(LayersTest, self).setUp() 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) self.layout_4d = dtensor.Layout.replicated(self.mesh, rank=4) self.layout_3d = dtensor.Layout.replicated(self.mesh, rank=3) self.layout_2d = dtensor.Layout.replicated(self.mesh, rank=2) self.layout_1d = dtensor.Layout.replicated(self.mesh, rank=1)
def setUp(self): super(LayoutMapTest, self).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) self.layout_2d = dtensor.Layout.replicated(self.mesh, rank=2) self.layout_1d = dtensor.Layout.replicated(self.mesh, rank=1) self.sharded_2d = dtensor.Layout.batch_sharded(self.mesh, 'X', rank=2) self.sharded_1d = dtensor.Layout.batch_sharded(self.mesh, 'X', rank=1)