class_docstr: An Asymmetric Cheeger Cut Pooling layer from the paper
  > [Total Variation Graph Neural Networks](https://arxiv.org/abs/2211.06218)<br>
  > Jonas Berg Hansen and 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 includes two auxiliary loss terms/components:
  A graph total variation component given by
  $$
      L_\text{GTV} = \frac{1}{2E} \sum_{k=1}^K \sum_{i=1}^N \sum_{j=i}^N a_{i,j} |s_{i,k} - s_{j,k}|,
  $$
  where \(E\) is the number of edges/links, \(K\) is the number of clusters or output nodes, and \(N\) is the number of nodes.
  
  An asymmetrical norm component given by
  $$
      L_\text{AN} = \frac{N(K - 1) - \sum_{k=1}^K ||\s_{:,k} - \textrm{quant}_{K-1} (\s_{:,k})||_{1, K-1}}{N(K-1)},
  $$
  
  The layer can be used without a supervised loss to compute node clustering by
  minimizing the two auxiliary losses.
  
  **Input**
  
  - Node features of shape `(batch, n_nodes_in, n_node_features)`;
  - Adjacency matrix of shape `(batch, n_nodes_in, n_nodes_in)`;
  
  **Output**
  
  - Reduced node features of shape `(batch, n_nodes_out, n_node_features)`;
  - 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;
  - `totvar_coeff`: coefficient for graph total variation loss component;
  - `balance_coeff`: coefficient for asymmetric norm loss component;  
  - `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_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: ["An","Asymmetric","Cheeger","Cut","Pooling","layer","from","the","paper",">","[","Total","Variation","Graph","Neural","Networks","]","(","https",":","\/\/arxiv.org\/abs\/2211.06218",")","<","br",">",">","Jonas","Berg","Hansen","and","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","includes","two","auxiliary","loss","terms\/components",":","A","graph","total","variation","component","given","by","$","$","L_\\text","{","GTV","}","=","\\frac","{","1","}","{","2E","}","\\sum_","{","k=1","}","^K","\\sum_","{","i=1","}","^N","\\sum_","{","j=i","}","^N","a_","{","i",",","j","}","|s_","{","i",",","k","}","-","s_","{","j",",","k","}","|",",","$","$","where","\\","(","E\\",")","is","the","number","of","edges\/links",",","\\","(","K\\",")","is","the","number","of","clusters","or","output","nodes",",","and","\\","(","N\\",")","is","the","number","of","nodes",".","An","asymmetrical","norm","component","given","by","$","$","L_\\text","{","AN","}","=","\\frac","{","N","(","K","-","1",")","-","\\sum_","{","k=1","}","^K","||\\s_","{",":",",k","}","-","\\textrm","{","quant","}","_","{","K-1","}","(","\\s_","{",":",",k","}",")","||_","{","1",",","K-1","}","}","{","N","(","K-1",")","}",",","$","$","The","layer","can","be","used","without","a","supervised","loss","to","compute","node","clustering","by","minimizing","the","two","auxiliary","losses",".","*","*","Input","*","*","-","Node","features","of","shape","`","(","batch",",","n_nodes_in",",","n_node_features",")","`",";","-","Adjacency","matrix","of","shape","`","(","batch",",","n_nodes_in",",","n_nodes_in",")","`",";","*","*","Output","*","*","-","Reduced","node","features","of","shape","`","(","batch",",","n_nodes_out",",","n_node_features",")","`",";","-","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",";","-","`","totvar_coeff","`",":","coefficient","for","graph","total","variation","loss","component",";","-","`","balance_coeff","`",":","coefficient","for","asymmetric","norm","loss","component",";","-","`","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_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):
      n_nodes = tf.cast(tf.shape(s, out_type=tf.int32)[-2], s.dtype)
  
      # k-quantile
      idx = tf.cast(tf.math.floor(n_nodes / self.k) + 1, dtype=tf.int32)
      med = tf.math.top_k(tf.linalg.matrix_transpose(s), k=idx).values[
          ..., -1
      ]
      # Asymmetric l1-norm
      if K.ndim(s) == 2:
          loss = s - med
      else:
          loss = s - med[:, tf.newaxis, ...]
      loss = (tf.cast(loss >= 0, loss.dtype) * (self.k - 1) * loss) + (
          tf.cast(loss < 0, loss.dtype) * loss * -1.0
      )
      loss = tf.math.reduce_sum(loss, axis=(-2, -1))
      loss = (
          1 / (n_nodes * (self.k - 1)) * (n_nodes * (self.k - 1) - loss)
      )
  
      return loss
  
file_path: spektral/spektral/layers/pooling/asym_cheeger_cut_pool.py
imports: import tensorflow
  import tensorflow.keras.backend
  from tensorflow.keras import Sequential
  from tensorflow.keras.layers import Dense
  from spektral.layers import ops
  from spektral.layers.pooling.src import SRCPool
label_desc: Use image node_id 9 for calling the AsymCheegerCutPool obj's underlying member method code with example usage: obj.balance_loss(s) and returns: loss