def test_cell_shapes(self):
    """Check that all the NeuralStackCell tensor shapes are correct.
    """
    batch_size = 5
    embedding_size = 3
    memory_size = 6
    num_units = 8

    stack = neural_stack.NeuralStackCell(num_units, memory_size, embedding_size)
    stack.build(None)

    self.assertEqual([1, 1, memory_size, memory_size],
                     stack.get_read_mask(0).shape)

    stack_input = tf.zeros([batch_size, 1, embedding_size], dtype=tf.float32)
    zero_state = stack.zero_state(batch_size, tf.float32)
    (outputs, (stack_next_state)) = stack.call(stack_input, zero_state)

    # Make sure that stack output shapes match stack input shapes
    self.assertEqual(outputs.shape, stack_input.shape)

    assert_cell_shapes(self, stack_next_state, zero_state)
Check that all the NeuralStackCell tensor shapes are correct.
    
tensor2tensor/models/research/neural_stack_test.py