class_docstr: Decoder model with a language modeling head on top for ONNX Runtime inference.
class_docstr_tok: ["Decoder","model","with","a","language","modeling","head","on","top","for","ONNX","Runtime","inference","."]
code: def compute_past_key_values_output_shapes(
      self,
      input_ids: torch.Tensor,
      encoder_hidden_states: torch.Tensor,
      use_cache_branch: Optional[bool],
      past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
  ) -> Dict[str, int]:
      batch_size = input_ids.size(0)
  
      num_attention_heads = self.normalized_config.num_attention_heads
      embed_size_per_head = (
          self.normalized_config.hidden_size // num_attention_heads
      )
  
      sequence_length = input_ids.size(1)
      encoder_sequence_length = encoder_hidden_states.size(1)
      if past_key_values is not None and use_cache_branch is not False:
          # Here, use_cache_branch may be None in the case of separate decoder without/with past, or True if the with past branch
          # of a merged decoder is used
          sequence_length += past_key_values[0].size(2)
  
      self_attn_shape = (
          batch_size,
          num_attention_heads,
          sequence_length,
          embed_size_per_head,
      )
  
      if past_key_values is not None and use_cache_branch is True:
          cross_attn_shape = (
              0,
              num_attention_heads,
              1,
              embed_size_per_head,
          )
      else:
          cross_attn_shape = (
              batch_size,
              num_attention_heads,
              encoder_sequence_length,
              embed_size_per_head,
          )
  
      past_key_values_shapes = {}
      for idx, name in enumerate(self.key_value_output_names):
          is_self_attn = idx % 4 < 2
          # decoder with past does not ouput cross attention key/values as they are constants
          past_key_values_shapes[name] = (
              self_attn_shape
              if (is_self_attn or self.num_pkv == 2)
              else cross_attn_shape
          )
      return past_key_values_shapes
  
imports: from abc import abstractmethod
  from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple, Union
  import numpy
  import torch
  from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
  from onnxruntime import InferenceSession
  from ..utils import NormalizedConfigManager
  from ..utils.logging import warn_once
  from .utils import get_ordered_input_names, logging
label_desc: Use image node_id 2 for calling the ORTDecoderForSeq2Seq obj's underlying member method code with example usage: obj.compute_past_key_values_output_shapes(input_ids, encoder_hidden_states, use_cache_branch, past_key_values) and returns: past_key_values_shapes