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.","`",")"]
code: def __repr__(self) -> str:
      return (
          f"{self.__class__.__name__}({self.in_channels}, "
          f"{self.out_channels}, heads={self.heads}, "
          f"num_clusters={self.num_clusters})"
      )
  
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 5 for calling the MemPooling obj's underlying member method code with example usage: obj.__repr__() and returns: self, str, self, str, self, str, self, str, self, str