code: def pre_attention(self, segment_number, query_antecedent,
                      memory_antecedent, bias):
      
      with tf.variable_scope(self.name + "/pre_attention", reuse=tf.AUTO_REUSE):
        assert memory_antecedent is None, "We only support language modeling"
        with tf.control_dependencies([
            tf.assert_greater_equal(self.batch_size, tf.size(segment_number))]):
          difference = self.batch_size - tf.size(segment_number)
          segment_number = tf.pad(segment_number, [[0, difference]])
          reset_op = self.reset(tf.reshape(tf.where(
              tf.less(segment_number, self.segment_number)), [-1]))
        memory_results = {}
        with tf.control_dependencies([reset_op]):
          with tf.control_dependencies([
              self.update_segment_number(segment_number)]):
            x = tf.pad(query_antecedent, [
                [0, difference], [0, 0], [0, 0]])
            access_logits, retrieved_mem = self.read(x)
        memory_results["x"] = x
        memory_results["access_logits"] = access_logits
        memory_results["retrieved_mem"] = retrieved_mem
        return memory_results, query_antecedent, memory_antecedent, bias
docs: Called prior to self-attention, to incorporate memory items.
  
      Args:
        segment_number: an integer Tensor with shape [batch]
        query_antecedent: a Tensor with shape [batch, length_q, channels]
        memory_antecedent: must be None. Attention normally allows this to be a
          Tensor with shape [batch, length_m, channels], but we currently only
          support memory for decoder-side self-attention.
        bias: bias Tensor (see attention_bias())
      Returns:
        (data, new_query_antecedent, new_memory_antecedent, new_bias)