code: def attention(self, query, key, value, mask=None):
      score = torch.matmul(query, key.permute(0, 1, 3, 2))
      dim_key = torch.tensor(key.shape[-1]).type(torch.float32)
      scaled_score = score / torch.sqrt(dim_key)
      if mask:
          scaled_score = mask * scaled_score
      weights = F.softmax(scaled_score, dim=-1)
      output = torch.matmul(weights, value)
      return output, weights
  
imports: import logging
  import torch
  from torch import nn
  from torch.nn import functional
  from ludwig.utils.torch_utils import get_activation, LudwigModule
label_desc: Use image node_id 2 for calling the MultiHeadSelfAttention obj's underlying member method code with example usage: obj.attention(query, key, value, mask) and returns: output,  weights