code: def schedule_triton_mha(in_shape, fwd=True):
      # default
      ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 0]
      if in_shape == [256, 4, 256, 16]:
          ret = [64, 32, 2, 4] if fwd else [64, 64, 4, 0]
      elif in_shape == [128, 4, 256, 16]:
          ret = [64, 32, 2, 4] if fwd else [64, 64, 4, 0]
      elif in_shape == [64, 4, 256, 16]:
          ret = [64, 32, 2, 4] if fwd else [64, 64, 4, 0]
      elif in_shape == [32, 4, 256, 16]:
          ret = [64, 32, 2, 4] if fwd else [64, 64, 4, 0]
      # [*, 8, 256, 32]
      elif in_shape == [128, 8, 256, 32]:  # DAP1
          ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 1]
      elif in_shape == [64, 8, 256, 32]:  # DAP2
          ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 1]
      elif in_shape == [32, 8, 256, 32]:  # DAP4
          ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 1]
      elif in_shape == [16, 8, 256, 32]:  # DAP8
          ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 1]
      # [*, 8, 128, 32]
      elif in_shape == [256, 8, 128, 32]:  # DAP1
          ret = [64, 64, 4, 3] if fwd else [128, 64, 4, 1]
      elif in_shape == [128, 8, 128, 32]:  # DAP2
          ret = [128, 64, 4, 2] if fwd else [64, 64, 2, 0]
      elif in_shape == [64, 8, 128, 32]:  # DAP4
          ret = [128, 64, 4, 2] if fwd else [64, 64, 2, 0]
      elif in_shape == [32, 8, 128, 32]:  # DAP8
          ret = [128, 64, 4, 2] if fwd else [64, 64, 2, 0]
      # [*, 4, 256, 32]
      elif in_shape == [256, 4, 256, 32]:  # DAP1
          ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 0]
      elif in_shape == [128, 4, 256, 32]:  # DAP2
          ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 1]
      elif in_shape == [64, 4, 256, 32]:  # DAP4
          ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 1]
      elif in_shape == [32, 4, 256, 32]:  # DAP8
          ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 0]
      return ret[0], ret[1], ret[2], ret[3]
  
file_path: apex/apex/contrib/openfold_triton/mha.py
imports: import math
  from typing import Optional
  import torch
  import triton
  from einops import rearrange
  from apex.contrib.openfold_triton._mha_kernel import _attention_core, _bwd_kernel, _bwd_preprocess
label_desc: Use image node_id 5 for calling a global function with example usage: schedule_triton_mha(in_shape, fwd) and returns: ret,  ret,  ret,  ret