docs: def call(self, x, y, bias, cache=None):
      """Apply attention mechanism to x and y.
  
      Args:
        x: a tensor with shape [batch_size, length_x, hidden_size]
        y: a tensor with shape [batch_size, length_y, hidden_size]
        bias: attention bias that will be added to the result of the dot product.
        cache: (Used during prediction) dictionary with tensors containing results
          of previous attentions. The dictionary must have the items:
              {"k": tensor with shape [batch_size, i, key_channels],
               "v": tensor with shape [batch_size, i, value_channels]}
          where i is the current decoded length.
  
      Returns:
        Attention layer output with shape [batch_size, length_x, hidden_size]
      """
      # Linearly project the query (q), key (k) and value (v) using different
      # learned projections. This is in preparation of splitting them into
      # multiple heads. Multi-head attention uses multiple queries, keys, and
      # values rather than regular attention (which uses a single q, k, v).
      q = self.q_dense_layer(x)
      k = self.k_dense_layer(y)
      v = self.v_dense_layer(y)
  
      if cache is not None:
        # Combine cached keys and values with new keys and values.
        k = tf.concat([cache["k"], k], axis=1)
        v = tf.concat([cache["v"], v], axis=1)
  
        # Update cache
        cache["k"] = k
        cache["v"] = v
  
      # Split q, k, v into heads.
      q = self.split_heads(q)
      k = self.split_heads(k)
      v = self.split_heads(v)
  
      # Scale q to prevent the dot product between q and k from growing too large.
      depth = (self.hidden_size // self.num_heads)
      q *= depth ** -0.5
  
      # Calculate dot product attention
      logits = tf.matmul(q, k, transpose_b=True)
      logits += bias
      weights = tf.nn.softmax(logits, name="attention_weights")
      if self.train:
        mlperf_log.transformer_print(
            key=mlperf_log.MODEL_HP_ATTENTION_DROPOUT,
            value=self.attention_dropout)
        weights = tf.nn.dropout(weights, 1.0 - self.attention_dropout)
      attention_output = tf.matmul(weights, v)
  
      # Recombine heads --> [batch_size, length, hidden_size]
      attention_output = self.combine_heads(attention_output)
  
      # Run the combined outputs through another linear projection layer.
      attention_output = self.output_dense_layer(attention_output)
      return attention_output