text: import tensorflow as tf
  import tensorflow.keras as ks
  import tensorflow.keras.backend as ksb
  
  from kgcnn.layers.base import GraphBaseLayer
  
  # Order Matters: Sequence to sequence for sets
  # by Vinyals et al. 2016
  # https://arxiv.org/abs/1511.06391
  
  
  @tf.keras.utils.register_keras_serializable(package='kgcnn', name='PoolingSet2Set')
  class PoolingSet2Set(GraphBaseLayer):
      """Pooling Node or edge embeddings by the Set2Set encoder part from layer.
      This was first proposed by `NMPNN <http://arxiv.org/abs/1704.01212>`_ .
      The Reading to Memory has to be handled separately.
      Uses a keras LSTM layer for the updates.
      
      Args:
          channels (int): Number of channels for the LSTM update.
          T (int): Numer of iterations. Default is T=3.
          pooling_method : Pooling method for PoolingSet2Set. Default is 'mean'.
          init_qstar: How to generate the first q_star vector. Default is 'mean'.
          activation: Activation function to use.
              Default: hyperbolic tangent (`tanh`). If you pass `None`, no activation
              is applied (ie. "linear" activation: `a(x) = x`).
          recurrent_activation: Activation function to use for the recurrent step.
              Default: sigmoid (`sigmoid`). If you pass `None`, no activation is
              applied (ie. "linear" activation: `a(x) = x`).
          use_bias: Boolean (default `True`), whether the layer uses a bias vector.
          kernel_initializer: Initializer for the `kernel` weights matrix, used for
              the linear transformation of the inputs. Default: `glorot_uniform`.
              recurrent_initializer: Initializer for the `recurrent_kernel` weights
              matrix, used for the linear transformation of the recurrent state.
              Default: `orthogonal`.
          bias_initializer: Initializer for the bias vector. Default: `zeros`.
              unit_forget_bias: Boolean (default `True`). If True, add 1 to the bias of
              the forget gate at initialization. Setting it to true will also force
              `bias_initializer="zeros"`. This is recommended in [Jozefowicz et
              al.](http://www.jmlr.org/proceedings/papers/v37/jozefowicz15.pdf).
          kernel_regularizer: Regularizer function applied to the `kernel` weights
              matrix. Default: `None`.
          recurrent_regularizer: Regularizer function applied to the
              `recurrent_kernel` weights matrix. Default: `None`.
              bias_regularizer: Regularizer function applied to the bias vector. Default:
              `None`.
          activity_regularizer: Regularizer function applied to the output of the
              layer (its "activation"). Default: `None`.
          kernel_constraint: Constraint function applied to the `kernel` weights
              matrix. Default: `None`.
          recurrent_constraint: Constraint function applied to the `recurrent_kernel`
              weights matrix. Default: `None`.
          bias_constraint: Constraint function applied to the bias vector. Default:
              `None`.
          dropout: Float between 0 and 1. Fraction of the units to drop for the linear
              transformation of the inputs. Default: 0.
              recurrent_dropout: Float between 0 and 1. Fraction of the units to drop for
              the linear transformation of the recurrent state. Default: 0.
          return_sequences: Boolean. Whether to return the last output. in the output
              sequence, or the full sequence. Default: `False`.
          return_state: Boolean. Whether to return the last state in addition to the
              output. Default: `False`.
          go_backwards: Boolean (default `False`). If True, process the input sequence
              backwards and return the reversed sequence.
          stateful: Boolean (default `False`). If True, the last state for each sample
              at index i in a batch will be used as initial state for the sample of
              index i in the following batch.
          time_major: The shape format of the `inputs` and `outputs` tensors.
              If True, the inputs and outputs will be in shape
              `[timesteps, batch, feature]`, whereas in the False case, it will be
              `[batch, timesteps, feature]`. Using `time_major = True` is a bit more
              efficient because it avoids transposes at the beginning and end of the
              RNN calculation. However, most TensorFlow data is batch-major, so by
              default this function accepts input and emits output in batch-major
              form.
          unroll: Boolean (default `False`). If True, the network will be unrolled,
              else a symbolic loop will be used. Unrolling can speed-up a RNN, although
              it tends to be more memory-intensive. Unrolling is only suitable for short
              sequences.
      """
  
      def __init__(self,
                   # Args
                   channels,
                   T=3,
                   pooling_method='mean',
                   init_qstar='mean',
                   # Args for LSTM
                   activation="tanh",
                   recurrent_activation="sigmoid",
                   use_bias=True,
                   kernel_initializer="glorot_uniform",
                   recurrent_initializer="orthogonal",
                   bias_initializer="zeros",
                   unit_forget_bias=True,
                   kernel_regularizer=None,
                   recurrent_regularizer=None,
                   bias_regularizer=None,
                   activity_regularizer=None,
                   kernel_constraint=None,
                   recurrent_constraint=None,
                   bias_constraint=None,
                   dropout=0.0,
                   recurrent_dropout=0.0,
                   implementation=2,
                   return_sequences=False,  # Should not be changed here
                   return_state=False,  # Should not be changed here
                   go_backwards=False,  # Should not be changed here
                   stateful=False,
                   time_major=False,
                   unroll=False,
                   **kwargs):
          """Init layer."""
          super(PoolingSet2Set, self).__init__(**kwargs)
          # Number of Channels to use in LSTM
          self.channels = channels
          self.T = T  # Number of Iterations to work on memory
          self.pooling_method = pooling_method
          self.init_qstar = init_qstar
  
          if self.pooling_method == 'mean':
              self._pool = ksb.mean
          elif self.pooling_method == 'sum':
              self._pool = ksb.sum
          else:
              raise TypeError("ERROR:kgcnn: Unknown pooling, choose: 'mean', 'sum', ...")
  
          if self.init_qstar == 'mean':
              self.qstar0 = self.init_qstar_mean
          else:
              self.qstar0 = self.init_qstar_0
          # ...
  
          # LSTM Layer to run on m
          self.lay_lstm = ks.layers.LSTM(channels,
                                         activation=activation,
                                         recurrent_activation=recurrent_activation,
                                         use_bias=use_bias,
                                         kernel_initializer=kernel_initializer,
                                         recurrent_initializer=recurrent_initializer,
                                         bias_initializer=bias_initializer,
                                         unit_forget_bias=unit_forget_bias,
                                         kernel_regularizer=kernel_regularizer,
                                         recurrent_regularizer=recurrent_regularizer,
                                         bias_regularizer=bias_regularizer,
                                         activity_regularizer=activity_regularizer,
                                         kernel_constraint=kernel_constraint,
                                         recurrent_constraint=recurrent_constraint,
                                         bias_constraint=bias_constraint,
                                         dropout=dropout,
                                         recurrent_dropout=recurrent_dropout,
                                         implementation=implementation,
                                         return_sequences=return_sequences,
                                         return_state=return_state,
                                         go_backwards=go_backwards,
                                         stateful=stateful,
                                         time_major=time_major,
                                         unroll=unroll
                                         )
  
      def build(self, input_shape):
          """Build layer."""
          super(PoolingSet2Set, self).build(input_shape)
  
      def call(self, inputs, **kwargs):
          """Forward pass.
  
          Args:
              inputs: nodes
  
                  - nodes (tf.RaggedTensor): Embeddings to be encoded of shape (batch, [N], F)
          
          Returns:
              tf.Tensor: Pooled tensor q_star of shape (batch, 1, 2*channels)
          """
          self.assert_ragged_input_rank(inputs)
          x, batch_num, batch_index = inputs.values, inputs.row_lengths(), inputs.value_rowids()
  
          # Reading to memory removed here, is to be done by separately
          m = x  # (batch*None, feat)
  
          # Initialize q0 and r0
          qstar = self.qstar0(m, batch_index, batch_num)
  
          # start loop
          for i in range(0, self.T):
              q = self.lay_lstm(qstar)  # (batch,feat)
              qt = tf.repeat(q, batch_num, axis=0)  # (batch*num,feat)
              et = self.f_et(m, qt)  # (batch*num,)
              # get at = exp(et)/sum(et) with sum(et)
              at = ksb.exp(et - self.get_scale_per_sample(et, batch_index, batch_num))  # (batch*num,)
              norm = self.get_norm(at, batch_index, batch_num)  # (batch*num,)
              at = norm * at  # (batch*num,) x (batch*num,)
              # calculate rt
              at = ksb.expand_dims(at, axis=1)
              rt = m * at  # (batch*num,feat) x (batch*num,1)
              rt = tf.math.segment_sum(rt, batch_index)  # (batch,feat)
              # qstar = [q,r]
              qstar = ksb.concatenate([q, rt], axis=1)  # (batch,2*feat)
              qstar = ksb.expand_dims(qstar, axis=1)  # (batch,1,2*feat)
  
          return qstar
  
      def f_et(self, fm, fq):
          """Function to compute scalar from m and q. Can apply sum or mean etc.
          
          Args:
               fm (tf.Tensor): of shape (batch*N, feat)
               fq (tf.Tensor): of shape (batch*N, feat)
              
          Returns:
              tf.Tensor: et of shape (batch*N, )
          """
          fet = self._pool(fm * fq, axis=1)  # (batch*N, 1)
          return fet
  
      @staticmethod
      def get_scale_per_batch(x):
          """Get re-scaling for the batch."""
          return tf.keras.backend.max(x, axis=0, keepdims=True)
  
      @staticmethod
      def get_scale_per_sample(x, ind, rep):
          """Get re-scaling for the sample."""
          out = tf.math.segment_max(x, ind)  # (batch,)
          out = tf.repeat(out, rep)  # (batch*num,)
          return out
  
      @staticmethod
      def get_norm(x, ind, rep):
          """Compute Norm."""
          norm = tf.math.segment_sum(x, ind)  # (batch,)
          norm = tf.math.reciprocal_no_nan(norm)  # (batch,)
          norm = tf.repeat(norm, rep, axis=0)  # (batch*num,)
          return norm
  
      def init_qstar_0(self, m, batch_index, batch_num):
          """Initialize the q0 with zeros."""
          batch_shape = ksb.shape(batch_num)
          return tf.zeros((batch_shape[0], 1, 2 * self.channels))
  
      def init_qstar_mean(self, m, batch_index, batch_num):
          """Initialize the q0 with mean."""
          # use q0=avg(m) (or q0=0)
          # batch_shape = ksb.shape(batch_num)
          q = tf.math.segment_mean(m, batch_index)  # (batch,feat)
          # r0
          qt = tf.repeat(q, batch_num, axis=0)  # (batch*num,feat)
          et = self.f_et(m, qt)  # (batch*num,)
          # get at = exp(et)/sum(et) with sum(et)=norm
          at = ksb.exp(et - self.get_scale_per_sample(et, batch_index, batch_num))  # (batch*num,)
          norm = self.get_norm(at, batch_index, batch_num)  # (batch*num,)
          at = norm * at  # (batch*num,) x (batch*num,)
          # calculate rt
          at = ksb.expand_dims(at, axis=1)  # (batch*num,1)
          rt = m * at  # (batch*num,feat) x (batch*num,1)
          rt = tf.math.segment_sum(rt, batch_index)  # (batch,feat)
          # [q0,r0]
          qstar = ksb.concatenate([q, rt], axis=1)  # (batch,2*feat)
          qstar = ksb.expand_dims(qstar, axis=1)  # (batch,1,2*feat)
          return qstar
  
      def get_config(self):
          """Make config for layer."""
          config = super(PoolingSet2Set, self).get_config()
          config.update({"channels": self.channels, "T": self.T, "pooling_method": self.pooling_method,
                         "init_qstar": self.init_qstar})
          lstm_conf = self.lay_lstm.get_config()
          lstm_param = ["activation",
                        "recurrent_activation",
                        "use_bias",
                        "kernel_initializer",
                        "recurrent_initializer",
                        "bias_initializer",
                        "unit_forget_bias",
                        "kernel_regularizer",
                        "recurrent_regularizer",
                        "bias_regularizer",
                        "activity_regularizer",
                        "kernel_constraint",
                        "recurrent_constraint",
                        "bias_constraint",
                        "dropout",
                        "recurrent_dropout",
                        "implementation",
                        "return_sequences",  # Should not be changed here
                        "return_state",  # Should not be changed here
                        "go_backwards",  # Should not be changed here
                        "stateful",
                        "time_major",
                        "unroll"]
          for x in lstm_param:
              config.update({x: lstm_conf[x]})
          return config
  