class_docstr: A memory module that caches features for recent tokens.
  
  When the number of tokens cached is equal to the chunk size, this is
  equivalent to the memory used by Transformer-XL
  (https://arxiv.org/abs/1901.02860)
class_docstr_tok: ["A","memory","module","that","caches","features","for","recent","tokens",".","When","the","number","of","tokens","cached","is","equal","to","the","chunk","size",",","this","is","equivalent","to","the","memory","used","by","Transformer-XL","(","https",":","\/\/arxiv.org\/abs\/1901.02860",")"]
docstr: Called prior to self-attention, to incorporate memory items.
  
  Args:
    segment: 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",":","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, query_antecedent, memory_antecedent, bias
  ):
      """Called prior to self-attention, to incorporate memory items.
  
      Args:
        segment: 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)
      """
      assert (
          memory_antecedent is None
      ), "We only support language modeling"
  
      # In eval mode, batch size may be variable
      memory_batch_size = tf.shape(self.previous_vals)[0]
      current_batch_size = tf.shape(query_antecedent)[0]
      amount_to_pad = memory_batch_size - current_batch_size
  
      # If segment id is zero, don't attend back to the memory
      previous_bias = (
          self.previous_bias[:current_batch_size, :, :, :]
          + tf.cast(
              tf.equal(segment[:, None, None, None], 0), tf.float32
          )
          * -1e9
      )
  
      sliced_previous_vals = self.previous_vals[
          :current_batch_size, :, :
      ]
  
      new_memory_antecedent = tf.concat(
          [tf.stop_gradient(sliced_previous_vals), query_antecedent], 1
      )
      new_bias = tf.concat(
          [
              tf.tile(
                  tf.stop_gradient(previous_bias),
                  [1, 1, self.chunk_length, 1],
              ),
              tf.tile(bias, [current_batch_size, 1, 1, 1]),
          ],
          -1,
      )
  
      remember_segment = tf.pad(segment, [[0, amount_to_pad]])
      # TODO(kitaev): The code assumes that we always either increment the chunk
      # number or reset it to zero. This assumption will not hold if we re-run the
      # model for each token, e.g. for autoregressive greedy/beam/sampling decode.
      remember_vals = tf.pad(
          query_antecedent, [[0, amount_to_pad], [0, 0], [0, 0]]
      )
      # Query position is on axis -2 for bias: as long as a token can be attended
      # to from at least one query position (i.e. it's not padding), memorize it.
      remember_bias = tf.tile(
          tf.reduce_max(bias, -2, keepdims=True),
          [memory_batch_size, 1, 1, 1],
      )
      # Assume that query_antecedent is always a full chunk (i.e. not truncated)
      if self.chunk_length < self.tokens_to_cache:
          remember_vals = tf.concat(
              [self.previous_vals, remember_vals], 1
          )
          remember_bias = tf.concat(
              [
                  self.previous_bias
                  - 1e9
                  * tf.cast(
                      tf.equal(
                          tf.pad(segment, [[0, amount_to_pad]])[
                              :, None, None, None
                          ],
                          0,
                      ),
                      tf.float32,
                  ),
                  remember_bias,
              ],
              -1,
          )
      if self.chunk_length != self.tokens_to_cache:
          remember_vals = remember_vals[:, -self.tokens_to_cache :, :]
          remember_bias = remember_bias[
              :, :, :, -self.tokens_to_cache :
          ]
      token = (remember_segment, remember_vals, remember_bias)
  
      return token, query_antecedent, new_memory_antecedent, new_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 2 for calling the RecentTokensMemory obj's underlying member method code with example usage: obj.pre_attention(segment, query_antecedent, memory_antecedent, bias) and returns: token,  query_antecedent,  new_memory_antecedent,  new_bias