class_docstr: Memory based pooling layer from `"Memory-Based Graph Networks"
  <https://arxiv.org/abs/2002.09518>`_ paper, which learns a coarsened graph
  representation based on soft cluster assignments.
  
  .. math::
      S_{i,j}^{(h)} &= \frac{
      (1+{\| \mathbf{x}_i-\mathbf{k}^{(h)}_j \|}^2 / \tau)^{
      -\frac{1+\tau}{2}}}{
      \sum_{k=1}^K (1 + {\| \mathbf{x}_i-\mathbf{k}^{(h)}_k \|}^2 / \tau)^{
      -\frac{1+\tau}{2}}}
  
      \mathbf{S} &= \textrm{softmax}(\textrm{Conv2d}
      (\Vert_{h=1}^H \mathbf{S}^{(h)})) \in \mathbb{R}^{N \times K}
  
      \mathbf{X}^{\prime} &= \mathbf{S}^{\top} \mathbf{X} \mathbf{W} \in
      \mathbb{R}^{K \times F^{\prime}}
  
  where :math:`H` denotes the number of heads, and :math:`K` denotes the
  number of clusters.
  
  Args:
      in_channels (int): Size of each input sample :math:`F`.
      out_channels (int): Size of each output sample :math:`F^{\prime}`.
      heads (int): The number of heads :math:`H`.
      num_clusters (int): number of clusters :math:`K` per head.
      tau (int, optional): The temperature :math:`\tau`. (default: :obj:`1.`)
class_docstr_tok: ["Memory","based","pooling","layer","from","`","``","Memory-Based","Graph","Networks","''","<","https",":","\/\/arxiv.org\/abs\/2002.09518",">","`","_","paper",",","which","learns","a","coarsened","graph","representation","based","on","soft","cluster","assignments",".","..","math",":",":","S_","{","i",",","j","}","^","{","(","h",")","}","&","=","\\frac","{","(","1+","{","\\|","\\mathbf","{","x","}","_i-\\mathbf","{","k","}","^","{","(","h",")","}","_j","\\|","}","^2","\/","\\tau",")","^","{","-\\frac","{","1+\\tau","}","{","2","}","}","}","{","\\sum_","{","k=1","}","^K","(","1","+","{","\\|","\\mathbf","{","x","}","_i-\\mathbf","{","k","}","^","{","(","h",")","}","_k","\\|","}","^2","\/","\\tau",")","^","{","-\\frac","{","1+\\tau","}","{","2","}","}","}","\\mathbf","{","S","}","&","=","\\textrm","{","softmax","}","(","\\textrm","{","Conv2d","}","(","\\Vert_","{","h=1","}","^H","\\mathbf","{","S","}","^","{","(","h",")","}",")",")","\\in","\\mathbb","{","R","}","^","{","N","\\times","K","}","\\mathbf","{","X","}","^","{","\\prime","}","&","=","\\mathbf","{","S","}","^","{","\\top","}","\\mathbf","{","X","}","\\mathbf","{","W","}","\\in","\\mathbb","{","R","}","^","{","K","\\times","F^","{","\\prime","}","}","where",":","math",":","`","H","`","denotes","the","number","of","heads",",","and",":","math",":","`","K","`","denotes","the","number","of","clusters",".","Args",":","in_channels","(","int",")",":","Size","of","each","input","sample",":","math",":","`","F","`",".","out_channels","(","int",")",":","Size","of","each","output","sample",":","math",":","`","F^","{","\\prime","}","`",".","heads","(","int",")",":","The","number","of","heads",":","math",":","`","H","`",".","num_clusters","(","int",")",":","number","of","clusters",":","math",":","`","K","`","per","head",".","tau","(","int",",","optional",")",":","The","temperature",":","math",":","`","\\tau","`",".","(","default",":",":","obj",":","`","1.","`",")"]
docstr: Forward pass.
  
  Args:
      x (torch.Tensor): The node feature tensor of shape
          :math:`\mathbf{X} \in \mathbb{R}^{N \times F}` or
          :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`.
      batch (torch.Tensor, optional): The batch vector
          :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
          each node to a specific example.
          Should not be provided in case node features already have shape
          :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`.
          (default: :obj:`None`)
      mask (torch.Tensor, optional): A mask matrix
          :math:`\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}`, which
          indicates valid nodes for each graph when using
          node features of shape
          :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`.
          (default: :obj:`None`)
      max_num_nodes (int, optional): The size of the :math:`B` node
          dimension. Automatically calculated if not given.
          (default: :obj:`None`)
      batch_size (int, optional): The number of examples :math:`B`.
          Automatically calculated if not given. (default: :obj:`None`)
docstr_tok: ["Forward","pass",".","Args",":","x","(","torch.Tensor",")",":","The","node","feature","tensor","of","shape",":","math",":","`","\\mathbf","{","X","}","\\in","\\mathbb","{","R","}","^","{","N","\\times","F","}","`","or",":","math",":","`","\\mathbf","{","X","}","\\in","\\mathbb","{","R","}","^","{","B","\\times","N","\\times","F","}","`",".","batch","(","torch.Tensor",",","optional",")",":","The","batch","vector",":","math",":","`","\\mathbf","{","b","}","\\in","{","\\","{","0",",","\\ldots",",","B-1\\","}","}","^N","`",",","which","assigns","each","node","to","a","specific","example",".","Should","not","be","provided","in","case","node","features","already","have","shape",":","math",":","`","\\mathbf","{","X","}","\\in","\\mathbb","{","R","}","^","{","B","\\times","N","\\times","F","}","`",".","(","default",":",":","obj",":","`","None","`",")","mask","(","torch.Tensor",",","optional",")",":","A","mask","matrix",":","math",":","`","\\mathbf","{","M","}","\\in","{","\\","{","0",",","1","\\","}","}","^","{","B","\\times","N","}","`",",","which","indicates","valid","nodes","for","each","graph","when","using","node","features","of","shape",":","math",":","`","\\mathbf","{","X","}","\\in","\\mathbb","{","R","}","^","{","B","\\times","N","\\times","F","}","`",".","(","default",":",":","obj",":","`","None","`",")","max_num_nodes","(","int",",","optional",")",":","The","size","of","the",":","math",":","`","B","`","node","dimension",".","Automatically","calculated","if","not","given",".","(","default",":",":","obj",":","`","None","`",")","batch_size","(","int",",","optional",")",":","The","number","of","examples",":","math",":","`","B","`",".","Automatically","calculated","if","not","given",".","(","default",":",":","obj",":","`","None","`",")"]
code: def forward(
      self,
      x: Tensor,
      batch: Optional[Tensor] = None,
      mask: Optional[Tensor] = None,
      max_num_nodes: Optional[int] = None,
      batch_size: Optional[int] = None,
  ) -> Tuple[Tensor, Tensor]:
      r"""Forward pass.
  
      Args:
          x (torch.Tensor): The node feature tensor of shape
              :math:`\mathbf{X} \in \mathbb{R}^{N \times F}` or
              :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`.
          batch (torch.Tensor, optional): The batch vector
              :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
              each node to a specific example.
              Should not be provided in case node features already have shape
              :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`.
              (default: :obj:`None`)
          mask (torch.Tensor, optional): A mask matrix
              :math:`\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}`, which
              indicates valid nodes for each graph when using
              node features of shape
              :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`.
              (default: :obj:`None`)
          max_num_nodes (int, optional): The size of the :math:`B` node
              dimension. Automatically calculated if not given.
              (default: :obj:`None`)
          batch_size (int, optional): The number of examples :math:`B`.
              Automatically calculated if not given. (default: :obj:`None`)
      """
      if x.dim() <= 2:
          x, mask = to_dense_batch(
              x,
              batch,
              max_num_nodes=max_num_nodes,
              batch_size=batch_size,
          )
      elif mask is None:
          mask = x.new_ones((x.size(0), x.size(1)), dtype=torch.bool)
  
      (B, N, _), H, K = x.size(), self.heads, self.num_clusters
  
      dist = (
          torch.cdist(self.k.view(H * K, -1), x.view(B * N, -1), p=2)
          ** 2
      )
      dist = (1.0 + dist / self.tau).pow(-(self.tau + 1.0) / 2.0)
  
      dist = dist.view(H, K, B, N).permute(2, 0, 3, 1)  # [B, H, N, K]
      S = dist / dist.sum(dim=-1, keepdim=True)
  
      S = self.conv(S).squeeze(dim=1).softmax(dim=-1)  # [B, N, K]
      S = S * mask.view(B, N, 1)
  
      x = self.lin(S.transpose(1, 2) @ x)
  
      return x, S
  
file_path: pytorch_geometric/torch_geometric/nn/pool/mem_pool.py
imports: from typing import Optional, Tuple
  import torch
  from torch import Tensor
  from torch.nn import Conv2d, KLDivLoss, Linear, Parameter
  from torch_geometric.utils import to_dense_batch
label_desc: Use image node_id 4 for calling the MemPooling obj's underlying member method code with example usage: obj.forward(x, batch, mask, max_num_nodes, batch_size) and returns: x,  S