def test_parametric_attention_model_with_single_representation(self):

    model = parametric_attention.SimpleParametricAttention(
        output_dimension=2,
        input_embedding_dimension=2,
        vocab_size=10,
        num_representations=1,
        max_sequence_size=20)

    input_batch = tf.convert_to_tensor(
        np.random.randint(low=0, high=10, size=(10, 20)))
    output = model(input_batch)

    self.assertIsInstance(model, tf.keras.Model)
    self.assertSequenceEqual(output.numpy().shape, [10, 1, 2])