class_docstr: Attention layer with cache used for auto-agressive decoding.
  
  Attributes:
    num_heads: Number of attention heads.
    head_size: Size of each attention head.
    **kwargs: Other keyword arguments inherit from `Attention` class.
class_docstr_tok: ["Attention","layer","with","cache","used","for","auto-agressive","decoding",".","Attributes",":","num_heads",":","Number","of","attention","heads",".","head_size",":","Size","of","each","attention","head",".","*","*","kwargs",":","Other","keyword","arguments","inherit","from","`","Attention","`","class","."]
code: def call(self, inputs, decode_loop_step=None):
      from_tensor = inputs[0]
      to_tensor = inputs[1]
      attention_mask = inputs[2] if len(inputs) >= 3 else None
      cache = inputs[3] if len(inputs) >= 4 else None
      # Scalar dimensions referenced here:
      #   B = batch size (number of sequences)
      #   F = `from_tensor` sequence length
      #   T = `to_tensor` sequence length
      #   N = `num_attention_heads`
      #   H = `size_per_head`
      # `query_tensor` = [B, F, N ,H]
      query_tensor = self._query_dense(from_tensor)
  
      # `key_tensor` = [B, T, N, H]
      key_tensor = self._key_dense(to_tensor)
  
      # `value_tensor` = [B, T, N, H]
      value_tensor = self._value_dense(to_tensor)
  
      if cache:
          key_tensor, value_tensor = self._update_cache(
              key_tensor, value_tensor, cache, decode_loop_step
          )
  
      # Take the dot product between "query" and "key" to get the raw
      # attention scores.
      attention_scores = tf.einsum(
          "BTNH,BFNH->BNFT", key_tensor, query_tensor
      )
      attention_scores = tf.multiply(
          attention_scores, 1.0 / math.sqrt(float(self._head_size))
      )
  
      # Normalize the attention scores to probabilities.
      # `attention_probs` = [B, N, F, T]
      attention_probs = self._masked_softmax(
          [attention_scores, attention_mask]
      )
  
      # This is actually dropping out entire tokens to attend to, which might
      # seem a bit unusual, but is taken from the original Transformer paper.
      attention_probs = self._dropout(attention_probs)
  
      # `context_layer` = [B, F, N, H]
      return (
          tf.einsum("BNFT,BTNH->BFNH", attention_probs, value_tensor),
          cache,
      )
  
file_path: autodist/examples/benchmark/utils/modeling/layers/attention.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  import tensorflow
  from utils.modeling.layers import dense_einsum
  from utils.modeling.layers import masked_softmax
label_desc: Use image node_id 3 for calling the CachedAttention obj's underlying member method code with example usage: obj.call(inputs, decode_loop_step) and returns: tf,  cache