class_docstr: Implements the Memory module.
  
  Based on Neural Turing Machines: arXiv:1410.5401 [cs.NE]
class_docstr_tok: ["Implements","the","Memory","module",".","Based","on","Neural","Turing","Machines",":","arXiv:1410.5401","[","cs.NE","]"]
docstr: 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)
docstr_tok: ["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",")"]
code: def pre_attention(
      self, segment_number, query_antecedent, memory_antecedent, bias
  ):
      """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)
      """
      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,
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
label_desc: Use image node_id 10 for calling the TransformerMemory obj's underlying member method code with example usage: obj.pre_attention(segment_number, query_antecedent, memory_antecedent, bias) and returns: memory_results,  query_antecedent,  memory_antecedent,  bias