class_docstr: The Just Balance pooling layer from the paper
  
  > [Simplifying Clustering with Graph Neural Networks](https://arxiv.org/abs/2207.08779)<br>
  > Filippo Maria Bianchi
  
  **Mode**: single, batch.
  
  This layer learns a soft clustering of the input graph as follows:
  $$
  \begin{align}
      \S &= \textrm{MLP}(\X); \\
      \X' &= \S^\top \X \\
      \A' &= \S^\top \A \S; \\
  \end{align}
  $$
  where \(\textrm{MLP}\) is a multi-layer perceptron with softmax output.
  
  The layer adds the following auxiliary loss to the model
  $$
      L = - \mathrm{Tr}(\sqrt{ \S^\top \S })
  $$
  
  The layer can be used without a supervised loss to compute node clustering by
  minimizing the auxiliary loss.
  
  The layer is originally designed to be used in conjuction with a
  [GCNConv](https://graphneural.network/layers/convolution/#gcnconv)
  layer operating on the following connectivity matrix
  
  $$
      \tilde{\A} = \I - \delta (\I - \D^{-1/2} \A \D^{-1/2})
  $$
  
  **Input**
  
  - Node features of shape `(batch, n_nodes_in, n_node_features)`;
  - Connectivity 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;
  - `normalized_loss`: booelan, whether to normalize the auxiliary loss in [0,1];
  - `return_selection`: boolean, whether to return the selection matrix;
  - `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","Just","Balance","pooling","layer","from","the","paper",">","[","Simplifying","Clustering","with","Graph","Neural","Networks","]","(","https",":","\/\/arxiv.org\/abs\/2207.08779",")","<","br",">",">","Filippo","Maria","Bianchi","*","*","Mode","*","*",":","single",",","batch",".","This","layer","learns","a","soft","clustering","of","the","input","graph","as","follows",":","$","$","\\begin","{","align","}","\\S","&","=","\\textrm","{","MLP","}","(","\\X",")",";","\\\\","\\X","'","&","=","\\S^\\top","\\X","\\\\","\\A","'","&","=","\\S^\\top","\\A","\\S",";","\\\\","\\end","{","align","}","$","$","where","\\","(","\\textrm","{","MLP","}","\\",")","is","a","multi-layer","perceptron","with","softmax","output",".","The","layer","adds","the","following","auxiliary","loss","to","the","model","$","$","L","=","-","\\mathrm","{","Tr","}","(","\\sqrt","{","\\S^\\top","\\S","}",")","$","$","The","layer","can","be","used","without","a","supervised","loss","to","compute","node","clustering","by","minimizing","the","auxiliary","loss",".","The","layer","is","originally","designed","to","be","used","in","conjuction","with","a","[","GCNConv","]","(","https",":","\/\/graphneural.network\/layers\/convolution\/","#","gcnconv",")","layer","operating","on","the","following","connectivity","matrix","$","$","\\tilde","{","\\A","}","=","\\I","-","\\delta","(","\\I","-","\\D^","{","-1\/2","}","\\A","\\D^","{","-1\/2","}",")","$","$","*","*","Input","*","*","-","Node","features","of","shape","`","(","batch",",","n_nodes_in",",","n_node_features",")","`",";","-","Connectivity","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",";","-","`","normalized_loss","`",":","booelan",",","whether","to","normalize","the","auxiliary","loss","in","[","0,1","]",";","-","`","return_selection","`",":","boolean",",","whether","to","return","the","selection","matrix",";","-","`","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 balance_loss(self, s):
      ss = ops.modal_dot(s, s, transpose_a=True)
      loss = -tf.linalg.trace(tf.math.sqrt(ss))
  
      if self.normalized_loss:
          n = float(tf.shape(s, out_type=tf.int32)[-2])
          c = float(tf.shape(s, out_type=tf.int32)[-1])
          loss = loss / tf.math.sqrt(n * c)
      return loss
  
file_path: spektral/spektral/layers/pooling/just_balance_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 8 for calling the JustBalancePool obj's underlying member method code with example usage: obj.balance_loss(s) and returns: loss