class_docstr: The DMoN pooling layer from the paper
  
  > [Graph Clustering with Graph Neural Networks](https://arxiv.org/abs/2006.16904)<br>
  > Anton Tsitsulin et al.
  
  **Mode**: single, batch.
  
  This layer learns a soft clustering of the input graph as follows:
  $$
  \begin{align}
      \C &= \textrm{MLP}(\X); \\
      \X' &= \C^\top \X \\
      \A' &= \C^\top \A \C; \\
  \end{align}
  $$
  where \(\textrm{MLP}\) is a multi-layer perceptron with softmax output.
  
  Two auxiliary loss terms are also added to the model: the modularity loss
  $$
      L_m = - \frac{1}{2m} \mathrm{Tr}(\C^\top \A \C - \C^\top \d^\top \d \C)
  $$
  and the collapse regularization loss
  $$
      L_c = \frac{\sqrt{k}}{n} \left\|
          \sum_i \C_i^\top
      \right\|_F -1.
  $$
  
  This layer is based on the original implementation found
  [here](https://github.com/google-research/google-research/blob/master/graph_embedding/dmon/dmon.py).
  
  **Input**
  
  - Node features of shape `(batch, n_nodes_in, n_node_features)`;
  - Symmetrically normalized adjacency matrix of shape
  `(batch, n_nodes_in, n_nodes_in)`;
  
  **Output**
  
  - Reduced node features of shape `(batch, n_nodes_out, n_node_features)`;
  - Reduced adjacency matrix of shape `(batch, n_nodes_out, n_nodes_out)`;
  - If `return_selection=True`, the selection matrix of shape
  `(batch, n_nodes_in, n_nodes_out)`.
  
  **Arguments**
  
  - `k`: number of output nodes;
  - `mlp_hidden`: list of integers, number of hidden units for each hidden layer in
  the MLP used to compute cluster assignments (if `None`, the MLP has only one output
  layer);
  - `mlp_activation`: activation for the MLP layers;
  - `collapse_regularization`: strength of the collapse regularization;
  - `return_selection`: boolean, whether to return the selection matrix;
  - `use_bias`: use bias in the MLP;
  - `kernel_initializer`: initializer for the weights of the MLP;
  - `bias_initializer`: initializer for the bias of the MLP;
  - `kernel_regularizer`: regularization applied to the weights of the MLP;
  - `bias_regularizer`: regularization applied to the bias of the MLP;
  - `kernel_constraint`: constraint applied to the weights of the MLP;
  - `bias_constraint`: constraint applied to the bias of the MLP;
class_docstr_tok: ["The","DMoN","pooling","layer","from","the","paper",">","[","Graph","Clustering","with","Graph","Neural","Networks","]","(","https",":","\/\/arxiv.org\/abs\/2006.16904",")","<","br",">",">","Anton","Tsitsulin","et","al",".","*","*","Mode","*","*",":","single",",","batch",".","This","layer","learns","a","soft","clustering","of","the","input","graph","as","follows",":","$","$","\\begin","{","align","}","\\C","&","=","\\textrm","{","MLP","}","(","\\X",")",";","\\\\","\\X","'","&","=","\\C^\\top","\\X","\\\\","\\A","'","&","=","\\C^\\top","\\A","\\C",";","\\\\","\\end","{","align","}","$","$","where","\\","(","\\textrm","{","MLP","}","\\",")","is","a","multi-layer","perceptron","with","softmax","output",".","Two","auxiliary","loss","terms","are","also","added","to","the","model",":","the","modularity","loss","$","$","L_m","=","-","\\frac","{","1","}","{","2m","}","\\mathrm","{","Tr","}","(","\\C^\\top","\\A","\\C","-","\\C^\\top","\\d^\\top","\\d","\\C",")","$","$","and","the","collapse","regularization","loss","$","$","L_c","=","\\frac","{","\\sqrt","{","k","}","}","{","n","}","\\left\\|","\\sum_i","\\C_i^\\top","\\right\\|_F","-1",".","$","$","This","layer","is","based","on","the","original","implementation","found","[","here","]","(","https",":","\/\/github.com\/google-research\/google-research\/blob\/master\/graph_embedding\/dmon\/dmon.py",")",".","*","*","Input","*","*","-","Node","features","of","shape","`","(","batch",",","n_nodes_in",",","n_node_features",")","`",";","-","Symmetrically","normalized","adjacency","matrix","of","shape","`","(","batch",",","n_nodes_in",",","n_nodes_in",")","`",";","*","*","Output","*","*","-","Reduced","node","features","of","shape","`","(","batch",",","n_nodes_out",",","n_node_features",")","`",";","-","Reduced","adjacency","matrix","of","shape","`","(","batch",",","n_nodes_out",",","n_nodes_out",")","`",";","-","If","`","return_selection=True","`",",","the","selection","matrix","of","shape","`","(","batch",",","n_nodes_in",",","n_nodes_out",")","`",".","*","*","Arguments","*","*","-","`","k","`",":","number","of","output","nodes",";","-","`","mlp_hidden","`",":","list","of","integers",",","number","of","hidden","units","for","each","hidden","layer","in","the","MLP","used","to","compute","cluster","assignments","(","if","`","None","`",",","the","MLP","has","only","one","output","layer",")",";","-","`","mlp_activation","`",":","activation","for","the","MLP","layers",";","-","`","collapse_regularization","`",":","strength","of","the","collapse","regularization",";","-","`","return_selection","`",":","boolean",",","whether","to","return","the","selection","matrix",";","-","`","use_bias","`",":","use","bias","in","the","MLP",";","-","`","kernel_initializer","`",":","initializer","for","the","weights","of","the","MLP",";","-","`","bias_initializer","`",":","initializer","for","the","bias","of","the","MLP",";","-","`","kernel_regularizer","`",":","regularization","applied","to","the","weights","of","the","MLP",";","-","`","bias_regularizer","`",":","regularization","applied","to","the","bias","of","the","MLP",";","-","`","kernel_constraint","`",":","constraint","applied","to","the","weights","of","the","MLP",";","-","`","bias_constraint","`",":","constraint","applied","to","the","bias","of","the","MLP",";"]
code: def connect(self, a, s, **kwargs):
      a_pool = ops.matmul_at_b_a(s, a)
  
      # Modularity loss
      mod_loss = self.modularity_loss(a, s, a_pool)
      if K.ndim(a) == 3:
          mod_loss = K.mean(mod_loss)
      self.add_loss(mod_loss)
  
      return a_pool
  
file_path: spektral/spektral/layers/pooling/dmon_pool.py
imports: import tensorflow
  from tensorflow.keras import Sequential
  from tensorflow.keras import backend
  from tensorflow.keras.layers import Dense
  from spektral.layers import ops
  from spektral.layers.pooling.src import SRCPool
label_desc: Use image node_id 6 for calling the DMoNPool obj's underlying member method code with example usage: obj.connect(a, s) and returns: a_pool