class_docstr: Multi-head Attention block from Transformer models.
code: def complexity(cx, hidden_d, n_heads, seq_len):
      # See https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py
      h, w = cx["h"], cx["w"]
      flops, params, acts = cx["flops"], cx["params"], cx["acts"]
      # q, k, v = linear(input).chunk(3)
      flops += seq_len * (hidden_d * hidden_d * 3 + hidden_d * 3)
      params += hidden_d * hidden_d * 3 + hidden_d * 3
      acts += hidden_d * 3 * seq_len
      # attn_output_weights = torch.bmm(q, k.transpose)
      head_d = hidden_d // n_heads
      flops += n_heads * (seq_len * head_d * seq_len)
      acts += n_heads * seq_len * seq_len
      # attn_output = torch.bmm(attn_output_weights, v)
      flops += n_heads * (seq_len * seq_len * head_d)
      acts += n_heads * seq_len * head_d
      # attn_output = linear(attn_output)
      flops += seq_len * (hidden_d * hidden_d + hidden_d)
      params += hidden_d * hidden_d + hidden_d
      acts += hidden_d * seq_len
      return {
          "h": h,
          "w": w,
          "flops": flops,
          "params": params,
          "acts": acts,
      }
  
imports: import numpy
  import torch
  import torch.nn
  from pycls.core.config import cfg
  from torch.nn import Module
label_desc: Use image node_id 3 for calling the MultiheadAttention obj's underlying member method code with example usage: obj.complexity(cx, hidden_d, n_heads, seq_len) and returns: dict