text: import os
  from spektral.layers.convolutional.gcn_conv import GCNConv
  
  
  os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'
  import tensorflow as tf
  
  from spektral.layers import ECCConv
  from spektral.layers.pooling.global_pool import GlobalMaxPool, GlobalAvgPool, GlobalSumPool
  
  from tensorflow.keras import Model, Input
  from tensorflow.keras.layers import Dense, LeakyReLU, BatchNormalization
  from tensorflow.keras.activations import tanh
  from tensorflow.sparse import SparseTensor
  
  
  hidden_states = 64
  activation = LeakyReLU(alpha = 0.1)
  
  # Probably needs regularization, but first step is just to fit, then we will regularize.
  
  class model(Model):
      def __init__(self, n_out = 7):
          super().__init__()
          # Define layers of the model
          self.ECC1    = ECCConv(hidden_states, [hidden_states, hidden_states, hidden_states], n_out = hidden_states, activation = "relu")
          self.GCN1    = GCNConv(hidden_states, activation = "relu")
          self.GCN2    = GCNConv(hidden_states * 2, activation = "relu")
          self.GCN3    = GCNConv(hidden_states * 4, activation = "relu")
          self.GCN4    = GCNConv(hidden_states * 8, activation = "relu")
          self.Pool1   = GlobalMaxPool()
          self.Pool2   = GlobalAvgPool()
          self.Pool3   = GlobalSumPool()
          self.decode  = [Dense(size * hidden_states) for size in [16, 8, 4, 2, 2]]
          self.norm_layers  = [BatchNormalization() for i in range(len(self.decode))]
          self.d2      = Dense(n_out)
  
      def call(self, inputs, training = False):
          x, a, i = inputs
          a, e    = self.generate_edge_features(x, a)
          x = self.ECC1([x, a, e])
          x = self.GCN1([x, a])
          x = self.GCN2([x, a])
          x = self.GCN3([x, a])
          x = self.GCN4([x, a])
          x1 = self.Pool1([x, i])
          x2 = self.Pool2([x, i])
          x3 = self.Pool3([x, i])
          x = tf.concat([x1, x2, x3], axis = 1)
          for decode_layer, norm_layer in zip(self.decode, self.norm_layers):
            x = activation(decode_layer(x))
            x = norm_layer(x, training = training)
          x = self.d2(x)
          # tf.print(tf.shape(x))
          return x
  
      def generate_edge_features(self, x, a):
        send    = a.indices[:, 0]
        receive = a.indices[:, 1]
  
        forwards  = tf.gather(x[:, 3], send) <= tf.gather(x[:, 3], receive)
  
        send    = tf.cast(send[forwards], tf.int64)
        receive = tf.cast(receive[forwards], tf.int64)
  
        a       = SparseTensor(indices = tf.stack([send, receive], axis = 1), values = tf.ones(tf.shape(send), dtype = tf.float32), dense_shape = tf.cast(tf.shape(a), tf.int64))
  
        diff_x  = tf.subtract(tf.gather(x, receive), tf.gather(x, send))
  
        dists   = tf.sqrt(
          tf.reduce_sum(
            tf.square(
              diff_x[:, :3]
            ), axis = 1
          ))
  
        vects = tf.math.divide_no_nan(diff_x[:, :3], tf.expand_dims(dists, axis = -1))
  
        e = tf.concat([diff_x[:, 3:], tf.expand_dims(dists, -1), vects], axis = 1)
  
        return a, e
  