code: def replace_MatchedScaleMaskSoftmax(
      n: nn.Module,
  ) -> Optional[nn.Linear]:
      """
      Replaces MatchedScaleMaskSoftmax with exportable softmax layer
      Args:
          n: module to replace
      Returns:
          exportable module
      """
      # including the import here to avoid circular imports
      from nemo.collections.nlp.modules.common.megatron.fused_softmax import (
          MatchedScaleMaskSoftmax,
      )
  
      # disabling fusion for the MatchedScaleMaskSoftmax
      mod = MatchedScaleMaskSoftmax(
          n.input_in_fp16,
          n.input_in_bf16,
          n.attn_mask_type,
          False,
          n.mask_func,
          n.softmax_in_fp32,
          n.scale,
      )
      return mod
  
imports: import os
  from contextlib import nullcontext
  from enum import Enum
  from typing import Callable, Dict, Optional, Type
  import onnx
  import torch
  import torch.nn
  import torch.nn.functional
  from nemo.utils import CastToFloat, CastToFloatAll, logging
label_desc: Use image node_id 12 for calling a global function with example usage: replace_MatchedScaleMaskSoftmax(n) and returns: mod