code: def CanSchTriMHA(in_shape, has_bias=True, inf=1e9, training=True):
      if has_bias == False:  # skip bias is None
          return False
      if inf != 1e9:  # skip inf != 1e9
          return False
  
      lst_3d = in_shape[-3:]
      skip_neg2_dim = in_shape[:3] + in_shape[-1:]
      if not training and (
          in_shape == [1, 538, 4, 538, 16]
          or in_shape == [1, 585, 4, 585, 16]
          or in_shape == [1, 538, 4, 538, 32]
          or in_shape == [1, 585, 4, 585, 32]
          or in_shape == [1, 128, 8, 585, 32]
          or in_shape == [1, 128, 8, 538, 32]
          or lst_3d == [8, 128, 32]
          or skip_neg2_dim == [1, 1024, 8, 8]
          or skip_neg2_dim == [1, 128, 4, 32]
          or skip_neg2_dim == [1, 128, 8, 32]
      ):  # eval
          return False  # skip eval
      if (
          in_shape == [1, 256, 4, 256, 16]
          or in_shape == [1, 128, 4, 256, 16]
          or in_shape == [1, 64, 4, 256, 16]
          or in_shape == [1, 32, 4, 256, 16]
      ):  # 7.26%
          return True
      elif (
          in_shape == [1, 128, 8, 256, 32]
          or in_shape == [1, 64, 8, 256, 32]
          or in_shape == [1, 32, 8, 256, 32]
          or in_shape == [1, 16, 8, 256, 32]
      ):  # 21.77%
          return True
      elif (
          in_shape == [1, 256, 8, 128, 32]
          or in_shape == [1, 128, 8, 128, 32]
          or in_shape == [1, 64, 8, 128, 32]
          or in_shape == [1, 32, 8, 128, 32]
      ):  # 21.77% no bias
          return True
      elif (
          in_shape == [1, 256, 4, 256, 32]
          or in_shape == [1, 128, 4, 256, 32]
          or in_shape == [1, 64, 4, 256, 32]
          or in_shape == [1, 32, 4, 256, 32]
      ):  # 47.17%
          return True
      else:  # not support
          return False
  
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 4 for calling a global function with example usage: CanSchTriMHA(in_shape, has_bias, inf, training) and returns: False, False, False, True, True, True, True, False