def test_convert_mtf_tensor_to_np_array(self):
x_np = np.array([[1, 2, 3], [4, 5, 6]])
converter = test_utils.NumpyConverter()
shape = mtf.Shape([mtf.Dimension("dim0", 2), mtf.Dimension("dim1", 3)])
x_mtf = mtf.constant(converter.mesh, x_np, shape=shape, dtype=tf.int32)
actual = converter.convert_mtf_tensor_to_np_array(x_mtf)
self.assertAllEqual(x_np, actual)
def testConv1dValidPadding(self):
converter = test_utils.NumpyConverter()
batch = 2
d_model = 6
d_out = 1
length = 4
filter_size = 3
x = np.random.randn(batch, length, d_model)
x_mtf = converter.convert_np_array_to_mtf_tensor(
x, dtype=tf.float32, dim_names=["batch", "length", "d_model"])
conv_filter = np.random.randn(1, filter_size, d_model, d_out)
initializer = lambda shape, dtype, **kwargs: conv_filter
output_mtf = mtf.layers.conv1d(x_mtf,
output_dim=mtf.Dimension(
"output_dim", d_out),
filter_size=filter_size,
padding="VALID",
filter_initializer=initializer)
actual = converter.convert_mtf_tensor_to_np_array(output_mtf)
# Expected length is 2.
expected = np.empty(shape=(batch, 2, d_out), dtype=np.float32)
# [filter_size, d_model]
current_filter = conv_filter[0, :, :, 0]
# b: batch, k: filter_size, d: d_model.
expected[:, 0] = np.einsum("bkd,kd->b", x[:, :filter_size, :],
current_filter).reshape(batch, 1)
expected[:, 1] = np.einsum("bkd,kd->b", x[:, 1:, :],
current_filter).reshape(batch, 1)
self.assertAllClose(actual, expected)
def testReduceFirst(self):
converter = mtf_test_utils.NumpyConverter()
x = np.random.randint(0, 100, size=[2, 3, 6, 5])
x_mtf = converter.convert_np_array_to_mtf_tensor(x)
expected = x[:, :, 0, :]
reduced_mtf = mtf.reduce_first(x_mtf, reduced_dim=x_mtf.shape.dims[2])
actual = converter.convert_mtf_tensor_to_np_array(reduced_mtf)
self.assertAllEqual(expected, actual)
def test_layer_stack_call_pooled_length(self):
converter = test_utils.NumpyConverter()
x = np.random.randn(2, 3, 4, 5)
layer_stack, context = create_dummy_model(converter.mesh,
shapes=x.shape)
x_mtf = converter.convert_np_array_to_mtf_tensor(
x, dim_names=self.default_dim_names, dtype=tf.float32)
output_mtf = layer_stack.call(context, x_mtf)
actual = converter.convert_mtf_tensor_to_np_array(output_mtf)
self.assertAllEqual(actual.shape, (2, 3, 2, 5))
def testStrideTensor1d(self, pool_size):
converter = mtf_test_utils.NumpyConverter()
x = np.random.randint(0, 100, size=[2, 3, 6, 5])
x_mtf = converter.convert_np_array_to_mtf_tensor(x)
expected = x[:, :, range(0, x.shape[2], pool_size), :]
strided_mtf = mtf.stride_tensor_1d(x_mtf,
pool_dim=x_mtf.shape.dims[2],
pool_size=pool_size)
actual = converter.convert_mtf_tensor_to_np_array(strided_mtf)
self.assertAllEqual(expected, actual)
def test_convert_mtf_tensor_to_np_array_with_trainable_variable(self):
converter = test_utils.NumpyConverter()
shape = mtf.Shape([mtf.Dimension("dim0", 2), mtf.Dimension("dim1", 3)])
x_mtf = mtf.get_variable(converter.mesh,
name="x",
shape=shape,
dtype=tf.float32,
initializer=tf.zeros_initializer())
actual = converter.convert_mtf_tensor_to_np_array(x_mtf)
self.assertAllClose(np.zeros_like(actual), actual)
def testPoolTensor1d(self, pool_fn, pool_fn_mtf):
converter = mtf_test_utils.NumpyConverter()
pool_size = 2
x = np.random.randn(2, 3, 4, 5)
expected = np.empty(shape=[2, 3, 2, 5])
expected[:, :, 0, :] = pool_fn(x[:, :, 0:2, :], axis=2)
expected[:, :, 1, :] = pool_fn(x[:, :, 2:4, :], axis=2)
x_mtf = converter.convert_np_array_to_mtf_tensor(x, dtype=tf.float32)
pooled_mtf = mtf.pool_tensor_1d(x_mtf,
pool_dim=x_mtf.shape.dims[2],
reduce_fn=pool_fn_mtf,
pool_size=pool_size)
actual = converter.convert_mtf_tensor_to_np_array(pooled_mtf)
self.assertAllClose(expected, actual)
def test_conv1d_call_same_input_output_dims(self):
converter = test_utils.NumpyConverter()
batch = 2
d_model = 6
length = 3
inputs = np.random.randint(0, 10, size=[batch, length])
inputs_mtf = converter.convert_np_array_to_mtf_tensor(
inputs, dim_names=["batch", "length"])
# Dummy context with necessary information for Conv1DLayer.call
Context = collections.namedtuple("Context",
["inputs", "activation_dtype", "mode"])
context = Context(
inputs=inputs_mtf, activation_dtype=tf.float32, mode="train")
x = np.random.randn(batch, length, d_model)
x_mtf = converter.convert_np_array_to_mtf_tensor(
x, dtype=tf.float32, dim_names=["batch", "length", "d_model"])
conv_layer = transformer_layers.Conv1DLayer(
filter_size=3, output_size=d_model)
output_mtf = conv_layer.call(context, x_mtf)
self.assertAllEqual([batch, length, d_model],
output_mtf.shape.to_integer_list)
def test_layer_stack_call_padding_handling(self):
self.converter = test_utils.NumpyConverter()
x = np.random.randn(2, 3, 4, 5)
layer_stack, context = create_dummy_model(self.converter.mesh,
shapes=x.shape)
# The last two sequence positions are padding.
x[:, :, -2:, :] *= 0
sequence_id = np.ones_like(x, dtype=np.int32)
sequence_id[:, :, -2:, :] *= 0
context.sequence_id = self.converter.convert_np_array_to_mtf_tensor(
sequence_id, dim_names=self.default_dim_names)
x_mtf = self.converter.convert_np_array_to_mtf_tensor(
x, dim_names=self.default_dim_names, dtype=tf.float32)
output_mtf = layer_stack.call(context, x_mtf)
# [2, 3, 4, 5] -> [2, 3, 2, 5]
actual = self.converter.convert_mtf_tensor_to_np_array(output_mtf)
# After pooling, the last sequence position should be padding, i.e., zeros.
last_position = actual[:, :, -1, :]
self.assertAllClose(last_position, np.zeros_like(last_position))
def setUp(self):
super(TransformerLayersTest, self).setUp()
self.converter = test_utils.NumpyConverter()
def setUp(self):
super(FunnelTransformerTest, self).setUp()
self.converter = test_utils.NumpyConverter()
self.default_dim_names = ["outer_batch", "batch", "length", "d_model"]
