class_docstr: The Multihead Attention Block (MAB) from the `"Set Transformer: A
  Framework for Attention-based Permutation-Invariant Neural Networks"
  <https://arxiv.org/abs/1810.00825>`_ paper.
  
  .. math::
  
      \mathrm{MAB}(\mathbf{x}, \mathbf{y}) &= \mathrm{LayerNorm}(\mathbf{h} +
      \mathbf{W} \mathbf{h})
  
      \mathbf{h} &= \mathrm{LayerNorm}(\mathbf{x} +
      \mathrm{Multihead}(\mathbf{x}, \mathbf{y}, \mathbf{y}))
  
  Args:
      channels (int): Size of each input sample.
      heads (int, optional): Number of multi-head-attentions.
          (default: :obj:`1`)
      norm (str, optional): If set to :obj:`False`, will not apply layer
          normalization. (default: :obj:`True`)
      dropout (float, optional): Dropout probability of attention weights.
          (default: :obj:`0`)
class_docstr_tok: ["The","Multihead","Attention","Block","(","MAB",")","from","the","`","``","Set","Transformer",":","A","Framework","for","Attention-based","Permutation-Invariant","Neural","Networks","''","<","https",":","\/\/arxiv.org\/abs\/1810.00825",">","`","_","paper",".","..","math",":",":","\\mathrm","{","MAB","}","(","\\mathbf","{","x","}",",","\\mathbf","{","y","}",")","&","=","\\mathrm","{","LayerNorm","}","(","\\mathbf","{","h","}","+","\\mathbf","{","W","}","\\mathbf","{","h","}",")","\\mathbf","{","h","}","&","=","\\mathrm","{","LayerNorm","}","(","\\mathbf","{","x","}","+","\\mathrm","{","Multihead","}","(","\\mathbf","{","x","}",",","\\mathbf","{","y","}",",","\\mathbf","{","y","}",")",")","Args",":","channels","(","int",")",":","Size","of","each","input","sample",".","heads","(","int",",","optional",")",":","Number","of","multi-head-attentions",".","(","default",":",":","obj",":","`","1","`",")","norm","(","str",",","optional",")",":","If","set","to",":","obj",":","`","False","`",",","will","not","apply","layer","normalization",".","(","default",":",":","obj",":","`","True","`",")","dropout","(","float",",","optional",")",":","Dropout","probability","of","attention","weights",".","(","default",":",":","obj",":","`","0","`",")"]
code: def forward(
      self,
      x: Tensor,
      y: Tensor,
      x_mask: Optional[Tensor] = None,
      y_mask: Optional[Tensor] = None,
  ) -> Tensor:
      """"""  # noqa: D419
      if y_mask is not None:
          y_mask = ~y_mask
  
      out, _ = self.attn(x, y, y, y_mask, need_weights=False)
  
      if x_mask is not None:
          out[~x_mask] = 0.0
  
      out = out + x
  
      if self.layer_norm1 is not None:
          out = self.layer_norm1(out)
  
      out = out + self.lin(out).relu()
  
      if self.layer_norm2 is not None:
          out = self.layer_norm2(out)
  
      return out
  
imports: from typing import Optional
  import torch
  from torch import Tensor
  from torch.nn import LayerNorm, Linear, MultiheadAttention, Parameter
label_desc: Use image node_id 3 for calling the MultiheadAttentionBlock obj's underlying member method code with example usage: obj.forward(x, y, x_mask, y_mask) and returns: out