def state_size(self):
    """The NeuralStackCell maintains a tuple of state values.

    Returns:
      (controller_state.shape,
       read_values.shape,
       memory_values.shape,
       read_strengths.shape,
       write_strengths.shape)
    """
    return (tf.TensorShape([self._num_units]),
            tf.TensorShape([self._num_read_heads, self._embedding_size]),
            tf.TensorShape([self._memory_size, self._embedding_size]),
            tf.TensorShape([1, self._memory_size, 1]),
            tf.TensorShape([self._num_write_heads, self._memory_size, 1]))
The NeuralStackCell maintains a tuple of state values.

    Returns:
      (controller_state.shape,
       read_values.shape,
       memory_values.shape,
       read_strengths.shape,
       write_strengths.shape)
    
tensor2tensor/models/research/neural_stack.py