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 __init__(self, num_heads, head_size, **kwargs):
      super(CachedAttention, self).__init__(
          num_heads, head_size, **kwargs
      )
  
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 1 to create a new CachedAttention object from inherited base classes: Attention with example: obj = CachedAttention(num_heads, head_size)