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"]