old_contents: from megnet.layers import MEGNet
  from keras.layers import Input
  import unittest
  
  
  class TestLayer(unittest.TestCase):
      @classmethod
      def setUpClass(cls):
          cls.n_feature = 5
          cls.n_bond_features = 6
          cls.n_global_features = 2
          cls.x = [
              Input(shape=(None, cls.n_feature)),
              Input(shape=(None, cls.n_bond_features)),
              Input(shape=(None, cls.n_global_features)),
              Input(shape=(None, ), dtype='int32'),
              Input(shape=(None, ), dtype='int32'),
              Input(shape=(None, ), dtype='int32'),
              Input(shape=(None, ), dtype='int32'),
          ]
  
      def test_matgraph(self):
          units_v = [13, 14, 16]
          units_e = [16, 16, 17]
          units_u = [13, 14, 18]
          layer = MEGNet(units_v, units_e, units_u)
          out = layer(self.x)
          self.assertListEqual([i._keras_shape[-1] for i in out], [units_v[-1], units_e[-1], units_u[-1]])
          new_layer = MEGNet.from_config(layer.get_config())
          out2 = new_layer(self.x)
          self.assertListEqual([i._keras_shape[-1] for i in out2], [units_v[-1], units_e[-1], units_u[-1]])
  
  
  if __name__ == "__main__":
      unittest.main()
  
new_contents: from megnet.layers import MEGNet
  from keras.layers import Input
  import unittest
  
  
  class TestLayer(unittest.TestCase):
      @classmethod
      def setUpClass(cls):
          cls.n_feature = 5
          cls.n_bond_features = 6
          cls.n_global_features = 2
          cls.x = [
              Input(shape=(None, cls.n_feature)),
              Input(shape=(None, cls.n_bond_features)),
              Input(shape=(None, cls.n_global_features)),
              Input(shape=(None, ), dtype='int32'),
              Input(shape=(None, ), dtype='int32'),
              Input(shape=(None, ), dtype='int32'),
              Input(shape=(None, ), dtype='int32'),
          ]
  
      def test_megnet(self):
          units_v = [13, 14, 16]
          units_e = [16, 16, 17]
          units_u = [13, 14, 18]
          layer = MEGNet(units_v, units_e, units_u)
          out = layer(self.x)
          self.assertListEqual([i._keras_shape[-1] for i in out], [units_v[-1], units_e[-1], units_u[-1]])
          new_layer = MEGNet.from_config(layer.get_config())
          out2 = new_layer(self.x)
          self.assertListEqual([i._keras_shape[-1] for i in out2], [units_v[-1], units_e[-1], units_u[-1]])
  
  
  if __name__ == "__main__":
      unittest.main()
  
subject: change name to megnet