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)

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