code: def post_attention(self, token, x):
      
      with tf.variable_scope(self.name + "/post_attention", reuse=tf.AUTO_REUSE):
        depth = common_layers.shape_list(x)[-1]
        actual_batch_size = common_layers.shape_list(x)[0]
        memory_output = tf.gather(token["retrieved_mem"],
                                  tf.range(actual_batch_size))
        output = tf.add(tf.layers.dense(x, depth, use_bias=False),
                        tf.layers.dense(memory_output, depth))
        with tf.control_dependencies([output]):
          with tf.control_dependencies([
              self.write(token["x"], token["access_logits"])]):
            return tf.identity(output)
docs: Called after self-attention. The memory can be updated here.
  
      Args:
        token: Data returned by pre_attention, which can be used to carry over
          state related to the current memory operation.
        x: a Tensor of data after self-attention and feed-forward
      Returns:
        a (possibly modified) version of the input x