code: def _scaled_dot_product_attention(
      q: Tensor,
      k: Tensor,
      v: Tensor,
      attn_mask: Optional[Tensor] = None,
      dropout_p: float = 0.0,
      bsz: int = 1,
      subset_heads: Optional[Tensor] = None,
      subset_weights: Optional[Tensor] = None,
  ) -> Tuple[Tensor, Tensor]:
      B, Nt, E = q.shape
      q = q / math.sqrt(E)
      # B: bsz * total_num_heads
      # (B, Nt, E) x (B, E, Ns) -> (B, Nt, Ns)
      attn = torch.bmm(q, k.transpose(-2, -1))
      if attn_mask is not None:
          attn += attn_mask
      attn = softmax(attn, dim=-1)
      if dropout_p > 0.0:
          attn = dropout(attn, p=dropout_p)
      if subset_heads is None:
          # (B, Nt, Ns) x (B, Ns, E) -> (B, Nt, E)
          output = torch.bmm(attn, v)
      else:
          mixed_output = (
              torch.bmm(attn, v).contiguous().view(bsz, -1, Nt, E)
          )
          output = torch.stack(
              [
                  mixed_output[
                      torch.arange(bsz), subset_heads[:, col], :, :
                  ]
                  for col in range(subset_heads.size(1))
              ],
              dim=1,
          )
          output = output * subset_weights.unsqueeze(2).unsqueeze(3)
          output = output.contiguous().view(-1, Nt, E)
      if subset_heads is not None:
          _, Nt, Ns = attn.size()
          mixed_attn = attn.view(bsz, -1, Nt, Ns)
          attn = torch.stack(
              [
                  mixed_attn[
                      torch.arange(bsz), subset_heads[:, col], :, :
                  ]
                  for col in range(subset_heads.size(1))
              ],
              dim=1,
          )
      return output, attn
  
file_path: fairseq/examples/attention_head_selection/src/modules/multihead_functional.py
imports: from typing import Optional, Tuple
  import torch
  from torch import Tensor
  from torch.nn.functional import linear, softmax, dropout, pad, has_torch_function, handle_torch_function, _in_projection_packed
  import math
  import warnings
label_desc: Use image node_id 1 for calling a global function with example usage: _scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, bsz, subset_heads, subset_weights) and returns: output,  attn