class_docstr: A SortPool layer as described by
  [Zhang et al](https://www.cse.wustl.edu/~muhan/papers/AAAI_2018_DGCNN.pdf).
  This layers takes a graph signal \(\mathbf{X}\) and returns the topmost k
  rows according to the last column.
  If \(\mathbf{X}\) has less than k rows, the result is zero-padded to k.
  
  **Mode**: single, disjoint, batch.
  
  **Input**
  
  - Node features of shape `([batch], n_nodes, n_node_features)`;
  - Graph IDs of shape `(n_nodes, )` (only in disjoint mode);
  
  **Output**
  
  - Pooled node features of shape `(batch, k, n_node_features)` (if single mode, shape will
  be `(1, k, n_node_features)`).
  
  **Arguments**
  
  - `k`: integer, number of nodes to keep;
class_docstr_tok: ["A","SortPool","layer","as","described","by","[","Zhang","et","al","]","(","https",":","\/\/www.cse.wustl.edu\/~muhan\/papers\/AAAI_2018_DGCNN.pdf",")",".","This","layers","takes","a","graph","signal","\\","(","\\mathbf","{","X","}","\\",")","and","returns","the","topmost","k","rows","according","to","the","last","column",".","If","\\","(","\\mathbf","{","X","}","\\",")","has","less","than","k","rows",",","the","result","is","zero-padded","to","k.","*","*","Mode","*","*",":","single",",","disjoint",",","batch",".","*","*","Input","*","*","-","Node","features","of","shape","`","(","[","batch","]",",","n_nodes",",","n_node_features",")","`",";","-","Graph","IDs","of","shape","`","(","n_nodes",",",")","`","(","only","in","disjoint","mode",")",";","*","*","Output","*","*","-","Pooled","node","features","of","shape","`","(","batch",",","k",",","n_node_features",")","`","(","if","single","mode",",","shape","will","be","`","(","1",",","k",",","n_node_features",")","`",")",".","*","*","Arguments","*","*","-","`","k","`",":","integer",",","number","of","nodes","to","keep",";"]
code: def call(self, inputs):
      if self.data_mode == "disjoint":
          X, I = inputs
          X = ops.disjoint_signal_to_batch(X, I)
      else:
          X = inputs
          if self.data_mode == "single":
              X = tf.expand_dims(X, 0)
  
      N = tf.shape(X)[-2]
      sort_perm = tf.argsort(X[..., -1], direction="DESCENDING")
      X_sorted = tf.gather(X, sort_perm, axis=-2, batch_dims=1)
  
      def truncate():
          _X_out = X_sorted[..., : self.k, :]
          return _X_out
  
      def pad():
          padding = [[0, 0], [0, self.k - N], [0, 0]]
          _X_out = tf.pad(X_sorted, padding)
          return _X_out
  
      X_out = tf.cond(tf.less_equal(self.k, N), truncate, pad)
  
      if self.data_mode == "single":
          X_out = tf.squeeze(X_out, [0])
          X_out.set_shape((self.k, self.F))
      elif self.data_mode == "batch" or self.data_mode == "disjoint":
          X_out.set_shape((None, self.k, self.F))
  
      return X_out
  
imports: import tensorflow
  from tensorflow.keras import backend
  from tensorflow.keras import constraints, initializers, regularizers
  from tensorflow.keras.layers import Dense, Layer
  from spektral.layers import ops
label_desc: Use image node_id 3 for calling the SortPool obj's underlying member method code with example usage: obj.call(inputs) and returns: X_out, _X_out, _X_out