code: def compute_mhsa(q, k, v, scale_factor=1, mask=None):
      # resulted shape will be: [batch, heads, tokens, tokens]
  
      k = paddle.transpose(k, perm=[0, 1, 3, 2])
      scaled_dot_prod = paddle.matmul(q, k) * scale_factor
  
      if mask is not None:
          assert mask.shape == scaled_dot_prod.shape[2:]
          scaled_dot_prod = scaled_dot_prod.masked_fill(mask, -np.inf)
  
      attention = paddle.nn.functional.softmax(scaled_dot_prod, axis=-1)
      # calc result per head
  
      return paddle.matmul(attention, v)
  
imports: from typing import List, Tuple
  import numpy
  import paddle
  import paddle.nn
  from paddle import Tensor
  from medicalseg.cvlibs import manager
label_desc: Use image node_id 1 for calling a global function with example usage: compute_mhsa(q, k, v, scale_factor, mask) and returns: paddle