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)