class_docstr: This is DP-friendly implementation of nn.MultiheadAttention.
  For full reference see original module refer to
  :class:`torch.nn.MultiheadAttention`.
  
  Current implementation leverages pytorch modules as building blocks
  to allow DP engine to calculate per-sample gradients.
  This is in contrast with original implementation based on nn.functional.
class_docstr_tok: ["This","is","DP-friendly","implementation","of","nn.MultiheadAttention",".","For","full","reference","see","original","module","refer","to",":","class",":","`","torch.nn.MultiheadAttention","`",".","Current","implementation","leverages","pytorch","modules","as","building","blocks","to","allow","DP","engine","to","calculate","per-sample","gradients",".","This","is","in","contrast","with","original","implementation","based","on","nn.functional","."]
code: def forward(
      self,
      query,
      key,
      value,
      key_padding_mask=None,
      need_weights=True,
      attn_mask=None,
      is_causal=False,
  ):
      is_batched = query.dim() == 3
  
      assert is_batched == True, "The query must have a dimension of 3."
  
      r"""
          As per https://github.com/pytorch/opacus/issues/596, we have to include ``is_causal`` as a dummy parameter of the function,
          since it is used in the ``forward`` function of parent class ``nn.TransformerEncoderLayer``.
          """
      assert (
          is_causal == False
      ), "We currently do not support causal mask. Will fix it in the future."
  
      r"""
          Using the same logic with ``nn.MultiheadAttention`` (https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html).
          """
      if self.batch_first:
          if key is value:
              if query is key:
                  query = key = value = query.transpose(1, 0)
              else:
                  query, key = [x.transpose(1, 0) for x in (query, key)]
                  value = key
          else:
              query, key, value = [
                  x.transpose(1, 0) for x in (query, key, value)
              ]
  
      tgt_len, bsz, embed_dim = query.size()
      if embed_dim != self.embed_dim:
          raise ValueError(
              f"query has as size of {embed_dim} while the embedding"
              " size is {self.embed_dim}"
          )
  
      head_dim = embed_dim // self.num_heads
      if head_dim * self.num_heads != embed_dim:
          raise ValueError(
              f"embedding dimension {embed_dim} not divisible "
              "by number of heads {num_heads}"
          )
      scaling = float(head_dim) ** -0.5
  
      q = self.qlinear(query)
      k = self.klinear(key)
      v = self.vlinear(value)
  
      q = q * scaling
  
      if attn_mask is not None:
          if attn_mask.dtype not in (
              torch.float32,
              torch.float64,
              torch.uint8,
              torch.bool,
          ):
              raise ValueError(
                  f"Only float, byte, and bool types are supported for attn_mask, "
                  f"not {attn_mask.dtype}."
              )
  
          if attn_mask.dtype == torch.uint8:
              warnings.warn(
                  "Byte tensor for attn_mask in nn.MultiheadAttention is deprecated."
                  "Use bool tensor instead."
              )
              attn_mask = attn_mask.to(torch.bool)
  
          if attn_mask.dim() == 2:
              attn_mask = attn_mask.unsqueeze(0)
              if list(attn_mask.size()) != [
                  1,
                  query.size(0),
                  key.size(0),
              ]:
                  raise ValueError(
                      "The size of the 2D attn_mask is not correct."
                  )
          elif attn_mask.dim() == 3:
              if list(attn_mask.size()) != [
                  bsz * self.num_heads,
                  query.size(0),
                  key.size(0),
              ]:
                  raise ValueError(
                      "The size of the 3D attn_mask is not correct."
                  )
          else:
              raise ValueError(
                  "attn_mask's dimension {} is not supported".format(
                      attn_mask.dim()
                  )
              )
          # attn_mask's dim is 3 now.
  
      # convert ByteTensor key_padding_mask to bool
      if (
          key_padding_mask is not None
          and key_padding_mask.dtype == torch.uint8
      ):
          warnings.warn(
              "Byte tensor for key_padding_mask in nn.MultiheadAttention"
              "is deprecated. Use bool tensor instead."
          )
          key_padding_mask = key_padding_mask.to(torch.bool)
  
      if self.add_bias_kv:
          k = self.seq_bias_k(k)
          v = self.seq_bias_v(v)
          if attn_mask is not None:
              attn_mask = F.pad(attn_mask, (0, 1))
          if key_padding_mask is not None:
              key_padding_mask = F.pad(key_padding_mask, (0, 1))
  
      q = (
          q.contiguous()
          .view(tgt_len, bsz * self.num_heads, head_dim)
          .transpose(0, 1)
      )
      if k is not None:
          k = (
              k.contiguous()
              .view(-1, bsz * self.num_heads, head_dim)
              .transpose(0, 1)
          )
      if v is not None:
          v = (
              v.contiguous()
              .view(-1, bsz * self.num_heads, head_dim)
              .transpose(0, 1)
          )
  
      src_len = k.size(1)
  
      if key_padding_mask is not None:
          assert key_padding_mask.size(0) == bsz
          assert key_padding_mask.size(1) == src_len
  
      if self.add_zero_attn:
          src_len += 1
          k = torch.cat(
              [
                  k,
                  torch.zeros(
                      (k.size(0), 1) + k.size()[2:],
                      dtype=k.dtype,
                      device=k.device,
                  ),
              ],
              dim=1,
          )
          v = torch.cat(
              [
                  v,
                  torch.zeros(
                      (v.size(0), 1) + v.size()[2:],
                      dtype=v.dtype,
                      device=v.device,
                  ),
              ],
              dim=1,
          )
          if attn_mask is not None:
              attn_mask = F.pad(attn_mask, (0, 1))
          if key_padding_mask is not None:
              key_padding_mask = F.pad(key_padding_mask, (0, 1))
  
      attn_output_weights = torch.bmm(q, k.transpose(1, 2))
      assert list(attn_output_weights.size()) == [
          bsz * self.num_heads,
          tgt_len,
          src_len,
      ]
  
      if attn_mask is not None:
          if attn_mask.dtype == torch.bool:
              attn_output_weights.masked_fill_(attn_mask, float("-inf"))
          else:
              attn_output_weights += attn_mask
  
      if key_padding_mask is not None:
          attn_output_weights = attn_output_weights.view(
              bsz, self.num_heads, tgt_len, src_len
          )
          attn_output_weights = attn_output_weights.masked_fill(
              key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf")
          )
          attn_output_weights = attn_output_weights.view(
              bsz * self.num_heads, tgt_len, src_len
          )
  
      attn_output_weights = F.softmax(attn_output_weights, dim=-1)
      attn_output_weights = self.dropout(attn_output_weights)
  
      attn_output = torch.bmm(attn_output_weights, v)
      assert list(attn_output.size()) == [
          bsz * self.num_heads,
          tgt_len,
          head_dim,
      ]
      attn_output = (
          attn_output.transpose(0, 1)
          .contiguous()
          .view(tgt_len, bsz, embed_dim)
      )
      attn_output = self.out_proj(attn_output)
  
      if self.batch_first:
          attn_output = attn_output.transpose(1, 0)
  
      if need_weights:
          # average attention weights over heads
          attn_output_weights = attn_output_weights.view(
              bsz, self.num_heads, tgt_len, src_len
          )
          return (
              attn_output,
              attn_output_weights.sum(dim=1) / self.num_heads,
          )
      else:
          return attn_output, None
  
file_path: opacus/opacus/layers/dp_multihead_attention.py
imports: import warnings
  from collections import OrderedDict
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 3 for calling the DPMultiheadAttention obj's underlying member method code with example usage: obj.forward(query, key, value, key_padding_mask, need_weights, attn_mask, is_causal) and returns: attn_output,  unknown, attn_output,  None