def call(self, inputs, prev_state):
    """Evaluates one timestep of the current neural stack cell.

    See section 3.4 of Grefenstette et al., 2015.

    Args:
      inputs: The inputs to the neural stack cell should be a tf.float32 tensor
        with shape [batch_size, embedding_size]
      prev_state: The NeuralStackState from the previous timestep.

    Returns:
      A tuple of the output of the stack as well as the new NeuralStackState.
    """
    batch_size = tf.shape(inputs)[0]

    # Call the controller and get controller interface values.
    with tf.control_dependencies([prev_state.read_strengths]):
      controller_output = self.call_controller(
          inputs, prev_state.read_values, prev_state.controller_state,
          batch_size)

    # Always write input values to memory regardless of push strength.
    # See Equation-1 in Grefenstette et al., 2015.
    new_memory_values = prev_state.memory_values + tf.reduce_sum(
        tf.expand_dims(controller_output.write_values, axis=2) *
        prev_state.write_strengths,
        axis=1)

    # Attenuate the read strengths of existing memory values depending on the
    # current pop strength.
    # See Equation-2 in Grefenstette et al., 2015.
    new_read_strengths = prev_state.read_strengths
    for h in range(self._num_read_heads - 1, -1, -1):
      new_read_strengths = tf.nn.relu(new_read_strengths - tf.nn.relu(
          tf.slice(controller_output.pop_strengths,
                   [0, h, 0, 0],
                   [-1, 1, -1, -1]) -
          tf.expand_dims(
              tf.reduce_sum(new_read_strengths * self.get_read_mask(h), axis=2),
              axis=3)))

    # Combine all write heads and their associated push values into a single set
    # of read weights.
    new_read_strengths += tf.reduce_sum(
        controller_output.push_strengths * prev_state.write_strengths,
        axis=1, keep_dims=True)

    # Calculate the "top" value of the stack by looking at read strengths.
    # See Equation-3 in Grefenstette et al., 2015.
    new_read_values = tf.reduce_sum(
        tf.minimum(
            new_read_strengths,
            tf.nn.relu(1 - tf.expand_dims(
                tf.reduce_sum(
                    new_read_strengths * tf.concat([
                        self.get_read_mask(h)
                        for h in range(self._num_read_heads)
                    ], axis=1),
                    axis=2),
                axis=3))
        ) * tf.expand_dims(new_memory_values, axis=1),
        axis=2)

    # Temporarily split write strengths apart so they can be shifted in
    # different directions.
    write_strengths_by_head = tf.split(prev_state.write_strengths,
                                       self._num_write_heads,
                                       axis=1)
    # Shift the write strengths for each write head in the direction indicated
    # by get_write_head_offset().
    new_write_strengths = tf.concat([
        tf.roll(write_strength, shift=self.get_write_head_offset(h), axis=2)
        for h, write_strength in enumerate(write_strengths_by_head)
    ], axis=1)

    return (controller_output.outputs, NeuralStackState(
        controller_state=controller_output.state,
        read_values=new_read_values,
        memory_values=new_memory_values,
        read_strengths=new_read_strengths,
        write_strengths=new_write_strengths))
Evaluates one timestep of the current neural stack cell.

    See section 3.4 of Grefenstette et al., 2015.

    Args:
      inputs: The inputs to the neural stack cell should be a tf.float32 tensor
        with shape [batch_size, embedding_size]
      prev_state: The NeuralStackState from the previous timestep.

    Returns:
      A tuple of the output of the stack as well as the new NeuralStackState.
    
tensor2tensor/models/research/neural_stack.py