code: def py_mha(qkv, amask, b, s, h, d):
      qkv = qkv.view(b, s, h, 3, d)
      q = qkv[:, :, :, 0, :].permute(0, 2, 1, 3)
      k = qkv[:, :, :, 1, :].permute(0, 2, 1, 3)
      v = qkv[:, :, :, 2, :].permute(0, 2, 1, 3)
      p = torch.matmul(q.float(), k.permute(0, 1, 3, 2).float())
      p_masked = p / math.sqrt(d) + (1.0 - amask) * -10000.0
      s = torch.softmax(p_masked, -1).to(qkv.dtype)
      ctx = torch.matmul(s, v)
      ctx = ctx.permute(0, 2, 1, 3).contiguous()
  
      ctx.retain_grad()
  
      return ctx
  
file_path: apex/apex/contrib/test/fmha/test_fmha.py
imports: import math
  import unittest
  import torch
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: py_mha(qkv, amask, b, s, h, d) and returns: ctx