def get_controller_shape(self, batch_size):
    """Define the output shapes of the neural stack controller.

    Making this a separate functions so that it can be used in unit tests.

    Args:
      batch_size: The size of the current batch of data.

    Returns:
      A tuple of shapes for each output returned from the controller.
    """
    return (
        # push_strengths,
        [batch_size, self._num_write_heads, 1, 1],
        # pop_strengths
        [batch_size, self._num_write_heads, 1, 1],
        # write_values
        [batch_size, self._num_write_heads, self._embedding_size],
        # outputs
        [batch_size, 1, self._embedding_size],
        # state
        [batch_size, self._num_units])
Define the output shapes of the neural stack controller.

    Making this a separate functions so that it can be used in unit tests.

    Args:
      batch_size: The size of the current batch of data.

    Returns:
      A tuple of shapes for each output returned from the controller.
    
tensor2tensor/models/research/neural_stack.py