def test_random_value_initializer(self, initializer_cls, init_args):
layout = dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED],
self.mesh)
shape = (4, 4)
initializer = initializer_cls(**init_args)
# Make sure to raise error when keras global seed is not set.
with self.assertRaisesRegex(ValueError, "set the global seed"):
initializer(shape=shape, layout=layout)
try:
tf_utils.set_random_seed(1337)
value = initializer(shape=shape, layout=layout)
self.assertEqual(value.shape, shape)
fetched_layout = dtensor.fetch_layout(value)
self.assertEqual(layout, fetched_layout)
# Make sure when same seed is set again, the new initializer should
# generate same result
tf_utils.set_random_seed(1337)
initializer = initializer_cls(**init_args)
new_value = initializer(shape=shape, layout=layout)
self.assertAllClose(value, new_value)
finally:
# Unset the keras global generator so that it doesn't affect other
# tests that need to verify the existence of global generator.
backend._SEED_GENERATOR.generator = None
def test_static_value_initializer(self, initializer_cls, init_args):
layout = dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED],
self.mesh)
shape = (4, 4)
initializer = initializer_cls(**init_args)
value = initializer(shape=shape, layout=layout)
normal_tensor_value = initializer(shape=shape)
self.assertEqual(value.shape, shape)
fetched_layout = dtensor.fetch_layout(value)
self.assertEqual(layout, fetched_layout)
self.assertAllClose(value, normal_tensor_value)
def test_dense_layer_with_layout(self):
dense = layers.Dense(10,
kernel_layout=self.layout_2d,
bias_layout=self.layout_1d)
inputs = np.random.randint(size=[32, 8], low=0, high=4)
inputs = tf.constant(inputs, dtype=tf.float32)
d_inputs = dtensor.copy_to_mesh(
inputs, dtensor.Layout.replicated(self.mesh, rank=2))
output = dense(d_inputs)
self.assertIsInstance(dense.kernel, dtensor.DVariable)
self.assertIsInstance(dense.bias, dtensor.DVariable)
expected_layout = dtensor.Layout(
[dtensor.UNSHARDED, dtensor.UNSHARDED], self.mesh)
self.assertEqual(dtensor.fetch_layout(output), expected_layout)
# Make sure to produce same output when layout is not used
tf_utils.set_random_seed(1337)
dense_2 = layers.Dense(10)
output_2 = dense_2(inputs)
self.assertAllClose(output, output_2)