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","."]
docstr: Updates cache states and gets full-length key/value tensors.
docstr_tok: ["Updates","cache","states","and","gets","full-length","key\/value","tensors","."]
code: def _update_cache(
      self, key_tensor, value_tensor, cache, decode_loop_step
  ):
      """Updates cache states and gets full-length key/value tensors."""
      # Combines cached keys and values with new keys and values.
      if decode_loop_step is not None:
          # TPU special case.
          key_seq_dim = cache["key"].shape.as_list()[1]
          indices = tf.reshape(
              tf.one_hot(
                  decode_loop_step, key_seq_dim, dtype=key_tensor.dtype
              ),
              [1, key_seq_dim, 1, 1],
          )
          key_tensor = cache["key"] + key_tensor * indices
          value_seq_dim = cache["value"].shape.as_list()[1]
          indices = tf.reshape(
              tf.one_hot(
                  decode_loop_step,
                  value_seq_dim,
                  dtype=value_tensor.dtype,
              ),
              [1, value_seq_dim, 1, 1],
          )
          value_tensor = cache["value"] + value_tensor * indices
      else:
          key_tensor = tf.concat(
              [tf.cast(cache["key"], key_tensor.dtype), key_tensor],
              axis=1,
          )
          value_tensor = tf.concat(
              [
                  tf.cast(cache["value"], value_tensor.dtype),
                  value_tensor,
              ],
              axis=1,
          )
  
      # Update cache
      cache["key"] = key_tensor
      cache["value"] = value_tensor
  
      return key_tensor, value_tensor
  
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 2 for calling the CachedAttention obj's underlying member method code with example usage: obj._update_cache(key_tensor, value_tensor, cache, decode_loop_step) and returns: key_tensor,  value_tensor