code: def t5_extended_attention_mask(
      attention_mask_list: List[Tensor],
  ) -> List[Tensor]:
      def attn_mask_postprocess(attn_mask):
          # [b, 1, s, s]
          extended_attention_mask = attn_mask.unsqueeze(1)
          return extended_attention_mask
  
      return [
          attn_mask_postprocess(attn_mask)
          for attn_mask in attention_mask_list
      ]
  
imports: import logging
  from typing import List, Literal, Optional
  import torch
  from torch import Tensor
  from megatron.core import InferenceParams, parallel_state, tensor_parallel
  from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
  from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
  from megatron.core.models.common.language_module.language_module import LanguageModule
  from megatron.core.transformer.enums import AttnMaskType, ModelType
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.spec_utils import ModuleSpec
  from megatron.core.transformer.transformer_block import TransformerBlock
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.utils import make_tp_sharded_tensor_for_checkpoint
label_desc: Use image node_id 1 for calling a global function with example usage: t5_extended_attention_mask(attention_mask_list) and returns: unknown, extended_attention_mask