code: def _attention_no_bias(
      query: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
      mask: torch.Tensor,
      inf: float,
  ) -> torch.Tensor:
      # query:  [*, num_heads, Q, c_hidden]
      # key:    [*, num_heads, K, c_hidden]
      # value:  [*, num_heads, V, c_hidden]
      # mask:   Logit mask tensor broadcastable to [*, num_heads, Q, K]
      # bias:   Optional logit bias tensor broadcastable to [*, num_heads, Q, K]
      # inf:    Safe infinity value.
      # assuming K == V
  
      key = torch.swapdims(key, -2, -1)
      # key: [*, num_heads, c_hidden, K]
  
      scaling = 1.0 / math.sqrt(query.size(-1))
      a = torch.matmul(query * scaling, key)
      # a: [*, num_heads, Q, K]
  
      a += (mask - 1.0) * inf
      # a: [*, num_heads, Q, K]
  
      a = torch.softmax(a, dim=-1)
      # a: [*, num_heads, Q, K]
  
      a = torch.matmul(a, value)
      # a: [*, num_heads, Q, c_hidden]
  
      return a
  
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 7 for calling a global function with example usage: _attention_no_bias(query, key, value, mask, inf) and returns: a