def test_model_shapes(self):
    """Test a few of the important output shapes for NeuralStackModel.
    """
    batch_size = 100
    seq_length = 80
    embedding_size = 64
    vocab_size = 128

    hparams = neural_stack.neural_stack()
    problem_hparams = contrib.training().HParams()

    problem_hparams.add_hparam("modality", {
        "inputs": modalities.ModalityType.SYMBOL,
        "targets": modalities.ModalityType.SYMBOL,
    })
    problem_hparams.add_hparam("vocab_size", {
        "inputs": vocab_size,
        "targets": vocab_size,
    })
    model = neural_stack.NeuralStackModel(hparams,
                                          problem_hparams=problem_hparams)

    features = {
        "inputs": tf.ones([batch_size, seq_length, 1, 1],
                          dtype=tf.int32),
        "targets": tf.ones([batch_size, seq_length, 1, 1], dtype=tf.int32)
    }

    transformed_features = model.bottom(features)

    self.assertEqual([batch_size, seq_length, 1, embedding_size],
                     transformed_features["inputs"].shape)

    logits = model.body(transformed_features)

    self.assertEqual([batch_size, seq_length, 1, embedding_size], logits.shape)
Test a few of the important output shapes for NeuralStackModel.
    
tensor2tensor/models/research/neural_stack_test.py