code: def post_attention(self, token, x):
      
      with tf.control_dependencies([
          self.previous_segment.assign(token[0]),
          self.previous_vals.assign(token[1]),
          self.previous_bias.assign(token[2]),
          ]):
        return tf.identity(x)
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