class_docstr: A class containing all of the functions supporting generation, to be used as a mixin in [`TFPreTrainedModel`].
class_docstr_tok: ["A","class","containing","all","of","the","functions","supporting","generation",",","to","be","used","as","a","mixin","in","[","`","TFPreTrainedModel","`","]","."]
code: def _update_model_kwargs_for_xla_generation(
      self,
      model_outputs: ModelOutput,
      model_kwargs: Dict[str, Any],
      cur_len: int,
      max_length: int,
      batch_size: int,
      is_encoder_decoder: bool = False,
      batch_axis: int = 0,
  ):
      def _initialize_attention(
          model_kwargs, num_padding_values, is_encoder_decoder
      ):
          """initializes the appropriate attention mask -- encoder-decoder models use `decoder_attention_mask`"""
          if is_encoder_decoder:
              # One 1 for decoder_start_token_id, 0s for the currently-unfilled locations in the past tensor,
              # 1s for the actual input_ids
              decoder_attention_mask = tf.concat(
                  [
                      tf.ones((batch_size, 1), dtype=tf.int32),
                      tf.zeros(
                          (batch_size, num_padding_values),
                          dtype=tf.int32,
                      ),
                      tf.ones((batch_size, 1), dtype=tf.int32),
                  ],
                  axis=1,
              )
              mask = {"decoder_attention_mask": decoder_attention_mask}
          else:
              attention_mask = model_kwargs.pop("attention_mask")
              # 0s for the currently-unfilled locations in the past tensor, 1s for the actual input_ids
              attention_mask = tf.concat(
                  [
                      attention_mask,
                      tf.zeros(
                          (batch_size, num_padding_values),
                          dtype=attention_mask.dtype,
                      ),
                      tf.ones(
                          (batch_size, 1), dtype=attention_mask.dtype
                      ),
                  ],
                  axis=1,
              )
              mask = {"attention_mask": attention_mask}
          return mask
  
      def _update_attention(
          model_kwargs, new_past_index, is_encoder_decoder
      ):
          """updates the appropriate attention mask -- encoder-decoder models use `decoder_attention_mask`"""
          update_start = (
              tf.constant([0, 1], dtype=tf.int32) * new_past_index
          )
          if is_encoder_decoder:
              decoder_attention_mask = model_kwargs.pop(
                  "decoder_attention_mask"
              )
              decoder_attention_mask_update_slice = tf.ones(
                  (batch_size, 1), dtype=decoder_attention_mask.dtype
              )
              decoder_attention_mask = dynamic_update_slice(
                  decoder_attention_mask,
                  decoder_attention_mask_update_slice,
                  update_start,
              )
              mask = {"decoder_attention_mask": decoder_attention_mask}
          else:
              attention_mask = model_kwargs.pop("attention_mask")
              attention_mask_update_slice = tf.ones(
                  (batch_size, 1), dtype=attention_mask.dtype
              )
              attention_mask = dynamic_update_slice(
                  attention_mask,
                  attention_mask_update_slice,
                  update_start,
              )
              mask = {"attention_mask": attention_mask}
          return mask
  
      def _initialize_past(past, num_padding_values, batch_axis):
          """initialize past with zeros -- the structure depends on `batch_axis`"""
          if batch_axis == 0:
              padding_values = tf.constant(
                  [[0, 0], [0, 0], [0, num_padding_values], [0, 0]],
                  dtype=tf.int32,
              )
              new_past = ()
              for past_layer in past:
                  new_past_layer = list(past_layer)
                  for i in range(len(new_past_layer[:2])):
                      new_past_layer[i] = tf.pad(
                          past_layer[i], padding_values
                      )
                  new_past += (tuple(new_past_layer),)
          else:
              padding_values = tf.scatter_nd(
                  indices=[[3, 1]],
                  updates=[num_padding_values],
                  shape=(5, 2),
              )
              new_past = list(past)
              for i in range(len(past)):
                  new_past[i] = tf.pad(past[i], padding_values)
          return new_past
  
      def _update_past(past, new_past_index, batch_axis):
          if batch_axis == 0:
              slice_start_base = tf.constant([0, 0, 1, 0])
              new_past = ()
              for past_layer in past:
                  new_past_layer = list(past_layer)
                  for i in range(len(new_past_layer[:2])):
                      update_slice = past_layer[i][:, :, -1:]
                      # Write the last slice to the first open location in the padded past array
                      # and then truncate the last slice off the array
                      new_past_layer[i] = dynamic_update_slice(
                          past_layer[i][:, :, :-1],
                          update_slice,
                          slice_start_base * new_past_index,
                      )
                  new_past += (tuple(new_past_layer),)
          else:
              slice_start_base = tf.constant([0, 0, 0, 1, 0])
              new_past = [None for _ in range(len(past))]
              for i in range(len(past)):
                  update_slice = past[i][:, :, :, -1:]
                  # Write the last slice to the first open location in the padded past array
                  # and then truncate the last slice off the array
                  new_past[i] = dynamic_update_slice(
                      past[i][:, :, :, :-1],
                      update_slice,
                      slice_start_base * new_past_index,
                  )
          return new_past
  
      if "past_key_values" in model_outputs:
          past = model_outputs.past_key_values
      elif "mems" in model_outputs:
          past = model_outputs.mems
      elif "past_buckets_states" in model_outputs:
          past = model_outputs.past_buckets_states
      else:
          raise ValueError(
              f"No known past variable found in model outputs (model outputs keys: {list(model_outputs.keys())})"
          )
      is_past_initialized = model_kwargs.pop("past", None) is not None
  
      if not is_past_initialized:
          # The padded version of `past` has a length of `max_length - 1`, as `past` holds information relative to
          # previous autoregressive generation steps (step 0 has no past, step 1 has 1 past value, ..., the last step
          # has `max_length - 1` past values).
          num_padding_values = max_length - cur_len - 1
          mask = _initialize_attention(
              model_kwargs, num_padding_values, is_encoder_decoder
          )
          new_past = _initialize_past(
              past, num_padding_values, batch_axis
          )
      else:
          # The new index of past to be filled corresponds to the current length of the sequence, with two
          # subtractions: -1 because past holds information regarding previous generation steps (read comment above)
          # and -1 again because in an array the index is the length of the array minus 1.
          new_past_index = cur_len - 2
          mask = _update_attention(
              model_kwargs, new_past_index, is_encoder_decoder
          )
          new_past = _update_past(past, new_past_index, batch_axis)
  
      # sets the updated variables (mask and past)
      model_kwargs.update(mask)
      model_kwargs["past"] = tuple(new_past)
  
      return model_kwargs
  
imports: import inspect
  import warnings
  from dataclasses import dataclass
  from typing import Any, Dict, List, Optional, Tuple, Union
  import numpy
  import tensorflow
  from tensorflow.compiler.tf2xla.python.xla import dynamic_update_slice
  from .generation_tf_logits_process import TFForcedBOSTokenLogitsProcessor, TFForcedEOSTokenLogitsProcessor, TFForceTokensLogitsProcessor, TFLogitsProcessorList, TFMinLengthLogitsProcessor, TFNoBadWordsLogitsProcessor, TFNoRepeatNGramLogitsProcessor, TFRepetitionPenaltyLogitsProcessor, TFSuppressTokensAtBeginLogitsProcessor, TFSuppressTokensLogitsProcessor, TFTemperatureLogitsWarper, TFTopKLogitsWarper, TFTopPLogitsWarper
  from .models.auto import TF_MODEL_FOR_CAUSAL_LM_MAPPING, TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, TF_MODEL_FOR_VISION_2_SEQ_MAPPING
  from .tf_utils import shape_list, stable_softmax
  from .utils import ModelOutput, logging
label_desc: Use image node_id 18 for calling the TFGenerationMixin obj's underlying member method code with example usage: obj._update_model_kwargs_for_xla_generation(model_outputs, model_kwargs, cur_len, max_length, batch_size, is_encoder_decoder, batch_axis) and returns: model_kwargs, mask, mask, new_past, new_past