code: def create_mixer_cls(config, layer_idx=None):
      head_dim = getattr(
          config,
          "head_dim",
          config.hidden_size // config.num_attention_heads,
      )
      softmax_scale = (
          1.0 if not config.scale_attn_weights else head_dim ** (-0.5)
      )
      if config.scale_attn_by_inverse_layer_idx:
          assert layer_idx is not None
          softmax_scale /= float(layer_idx + 1)
      use_flash_attn = getattr(config, "use_flash_attn", False)
      fused_bias_fc = getattr(config, "fused_bias_fc", False)
      bf16 = getattr(config, "bf16", False)
      mixer_cls = partial(
          MHA,
          num_heads=config.num_attention_heads,
          dropout=config.attn_pdrop,
          softmax_scale=softmax_scale,
          causal=True,
          fused_bias_fc=fused_bias_fc,
          use_flash_attn=use_flash_attn,
          bf16=bf16,
      )
      return mixer_cls
  
imports: import math
  from functools import partial
  from collections import namedtuple
  from collections.abc import Sequence
  import torch
  import torch.nn
  import torch.nn.functional
  from transformers.models.gpt2.configuration_gpt2 import GPT2Config
  from src.models.modules.mha import MHA
  from src.models.modules.mlp import Mlp, FusedDenseGeluDense
  from src.models.modules.block import Block
  from src.models.modules.embedding import GPT2Embeddings
label_desc: Use image node_id 1 for calling a global function with example usage: create_mixer_cls(config, layer_idx) and returns: mixer_cls