molcraft 0.1.0a7__py3-none-any.whl → 0.1.0a9__py3-none-any.whl

molcraft/layers.py

@@ -274,20 +274,14 @@ class GraphConv(GraphLayer):
         units (int):
             Dimensionality of the output space.
         activation (keras.layers.Activation, str, None):
-            Activation function to use. If not specified, a linear activation (a(x) = x) is used.
-            Default to `None`.
+            Activation function to be accessed via `self.activation`, and used for the 
+            `message()` and `update()` methods, if not overriden. Default to `relu`.
         use_bias (bool):
-            Whether bias should be used in dense layers. Default to `True`.
+            Whether bias should be used in the dense layers. Default to `True`.
         normalize (bool, str):
-            Whether `LayerNormalization` should be applied to the final node feature output.
-            To use `BatchNormalization`, specify `batch_norm`. Default to `False`.
-        skip_connect (bool, str):
-            Whether node feature input should be added to the node feature output. 
-            If node feature input dim is not equal to `units` (node feature output dim), 
-            a projection layer will automatically project the residual before adding it 
-            to the output. To use weighted skip connection, 
-            specify `weighted`. The weight multiplied with the skip connection is a 
-            learnable scalar. Default to `True`.
+            Whether normalization should be applied to the final output. Default to `False`.
+        skip_connect (bool):
+            Whether node feature input should be added to the node feature output. Default to `True`.
         kernel_initializer (keras.initializers.Initializer, str):
             Initializer for the kernel weight matrix of the dense layers.
             Default to `glorot_uniform`.
@@ -314,10 +308,10 @@ class GraphConv(GraphLayer):
     def __init__(
         self, 
         units: int = None, 
-        activation: str | keras.layers.Activation | None = None,
+        activation: str | keras.layers.Activation | None = 'relu',
         use_bias: bool = True,
-        normalize: bool | str = False,
-        skip_connect: bool | str = True, 
+        normalize: bool = False,
+        skip_connect: bool = True,
         **kwargs
     ) -> None:
         super().__init__(use_bias=use_bias, **kwargs)
@@ -341,56 +335,56 @@ class GraphConv(GraphLayer):
     def units(self):
         return self._units 
     
+    @property
+    def activation(self):
+        return self._activation
+    
     def build(self, spec: tensors.GraphTensor.Spec) -> None:
         if not self.units:
             raise ValueError(
                 f'`self.units` needs to be a positive integer. Found: {self.units}.'
             )
         node_feature_dim = spec.node['feature'].shape[-1]
-        self._project_input_node_feature = (
+        self._project_residual = (
             self._skip_connect and (node_feature_dim != self.units)
         )
-        if self._project_input_node_feature:
-            warn(
+        if self._project_residual:
+            warnings.warn(
                 '`skip_connect` is set to `True`, but found incompatible dim ' 
                 'between input (node feature dim) and output (`self.units`). '
                 'Automatically applying a projection layer to residual to '
-                'match input and output. '
+                'match input and output. ',
+                stacklevel=2,
             )
-            self._residual_projection = self.get_dense(
-                self.units, name='residual_projection'
+            self._residual_dense = self.get_dense(
+                self.units, name='residual_dense'
             )
 
-        skip_connect = str(self._skip_connect).lower()
-        self._use_weighted_skip_connection = skip_connect.startswith('weight')
-        if self._use_weighted_skip_connection:
-            self._skip_connection_weight = self.add_weight(
-                name='skip_connection_weight', 
-                shape=(), 
-                initializer='ones',
-                trainable=True,
-            )
-
-        if self._normalize:
-            if str(self._normalize).lower().startswith('batch'):
-                self._output_norm = keras.layers.BatchNormalization(
-                    name='output_batch_norm'
-                )
-            else:
-                self._output_norm = keras.layers.LayerNormalization(
-                    name='output_layer_norm'
-                )
-
-        self._has_edge_feature = 'feature' in spec.edge 
+        self.has_edge_feature = 'feature' in spec.edge
+        self.has_node_coordinate = 'coordinate' in spec.node
 
         has_overridden_message = self.__class__.message != GraphConv.message 
         if not has_overridden_message:
-            self._message_dense = self.get_dense(self.units)
+            self._message_intermediate_dense = self.get_dense(self.units)
+            self._message_intermediate_activation = self.activation
+            self._message_final_dense = self.get_dense(self.units)
+
+        has_overridden_aggregate = self.__class__.message != GraphConv.aggregate 
+        if not has_overridden_aggregate:
+            pass
 
         has_overridden_update = self.__class__.update != GraphConv.update 
         if not has_overridden_update:
-            self._output_dense = self.get_dense(self.units)
-            self._output_activation = self._activation
+            self._update_intermediate_dense = self.get_dense(self.units)
+            self._update_intermediate_activation = self.activation
+            self._update_final_dense = self.get_dense(self.units)
+
+        if not self._normalize:
+            self._normalization = keras.layers.Identity()
+        elif str(self._normalize).lower().startswith('layer'):
+            self._normalization = keras.layers.LayerNormalization()
+        else:
+            self._normalization = keras.layers.BatchNormalization()
 
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         """Forward pass.
@@ -402,10 +396,10 @@ class GraphConv(GraphLayer):
                 A `GraphTensor` instance.
         """
         if self._skip_connect:
-            input_node_feature = tensor.node['feature']
-            if self._project_input_node_feature:
-                input_node_feature = self._residual_projection(input_node_feature)
- 
+            residual = tensor.node['feature']
+            if self._project_residual:
+                residual = self._residual_dense(residual)
+
         message = self.message(tensor)
         if not isinstance(message, tensors.GraphTensor):
             message = tensor.update({'edge': {'message': message}})
@@ -417,24 +411,24 @@ class GraphConv(GraphLayer):
             aggregate = tensor.update({'node': {'aggregate': aggregate}})
         elif not 'aggregate' in aggregate.node:
             raise ValueError('Could not find `aggregate` in `node` output.')
-
+            
         update = self.update(aggregate)
         if not isinstance(update, tensors.GraphTensor):
             update = tensor.update({'node': {'feature': update}})
         elif not 'feature' in update.node:
             raise ValueError('Could not find `feature` in `node` output.')
-            
-        updated_node_feature = update.node['feature']
+        
+        if update.node['feature'].shape[-1] != self.units:
+            raise ValueError('Updated node `feature` is not equal to `self.units`.')
+
+        feature = update.node['feature']
 
         if self._skip_connect:
-            if self._use_weighted_skip_connection:
-                input_node_feature *= self._skip_connection_weight
-            updated_node_feature += input_node_feature
-        
-        if self._normalize:
-            updated_node_feature = self._output_norm(updated_node_feature)
+            feature += residual 
+
+        feature = self._normalization(feature)
 
-        return update.update({'node': {'feature': updated_node_feature}})
+        return update.update({'node': {'feature': feature}})
 
     def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         """Compute messages.
@@ -445,24 +439,38 @@ class GraphConv(GraphLayer):
             tensor:
                 The inputted `GraphTensor` instance.
         """
-        if not self._has_edge_feature:
-            message_feature = tensor.gather('feature', 'source')
-        else:
-            message_feature = keras.ops.concatenate(
+        message = keras.ops.concatenate(
+            [
+                tensor.gather('feature', 'source'),
+                tensor.gather('feature', 'target'),
+            ],
+            axis=-1
+        )
+        if self.has_edge_feature:
+            message = keras.ops.concatenate(
                 [
-                    tensor.gather('feature', 'source'),
+                    message, 
                     tensor.edge['feature']
                 ], 
                 axis=-1
             )
-        message = self._message_dense(message_feature)
-        return tensor.update(
-            {
-                'edge': {
-                    'message': message
-                }
-            }
-        )
+        if self.has_node_coordinate:
+            euclidean_distance = ops.euclidean_distance(
+                tensor.gather('coordinate', 'target'),
+                tensor.gather('coordinate', 'source'),
+                axis=-1
+            )
+            message = keras.ops.concatenate(
+                [
+                    message,
+                    euclidean_distance
+                ],
+                axis=-1
+            )
+        message = self._message_intermediate_dense(message)
+        message = self._message_intermediate_activation(message)
+        message = self._message_final_dense(message)
+        return tensor.update({'edge': {'message': message}})
 
     def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         """Aggregates messages.
@@ -473,14 +481,16 @@ class GraphConv(GraphLayer):
             tensor:
                 A `GraphTensor` instance containing a message.
         """
+        previous = tensor.node['feature']
         aggregate = tensor.aggregate('message', mode='mean')
+        aggregate = keras.ops.concatenate([aggregate, previous], axis=-1)
         return tensor.update(
             {
                 'node': {
-                    'aggregate': aggregate,
+                    'aggregate': aggregate, 
                 },
                 'edge': {
-                    'message': None
+                    'message': None,
                 }
             }
         )
@@ -495,21 +505,16 @@ class GraphConv(GraphLayer):
                 A `GraphTensor` instance containing aggregated messages 
                 (updated node features).
         """
-        feature = keras.ops.concatenate(
-            [
-                tensor.node['aggregate'],
-                tensor.node['feature']
-            ],
-            axis=-1
-        )
-        update = self._output_dense(feature)
-        update = self._output_activation(update)
+        aggregate = tensor.node['aggregate']
+        node_feature = self._update_intermediate_dense(aggregate)
+        node_feature = self._update_intermediate_activation(node_feature)
+        node_feature = self._update_final_dense(node_feature)
         return tensor.update(
             {
                 'node': {
-                    'feature': update,
+                    'feature': node_feature,
                     'aggregate': None,
-                }
+                },
             }
         )
 
@@ -563,14 +568,16 @@ class GIConv(GraphConv):
         activation: keras.layers.Activation | str | None = 'relu',
         use_bias: bool = True,
         normalize: bool = False,
+        skip_connect: bool = True,
         update_edge_feature: bool = True,
         **kwargs,
     ):
         super().__init__(
             units=units, 
             activation=activation,
-            normalize=normalize, 
             use_bias=use_bias, 
+            normalize=normalize, 
+            skip_connect=skip_connect,
             **kwargs
         )
         self._update_edge_feature = update_edge_feature
@@ -585,16 +592,16 @@ class GIConv(GraphConv):
             trainable=True,
         )
 
-        self._has_edge_feature = 'feature' in spec.edge
-        if self._has_edge_feature:
+        if self.has_edge_feature:
             edge_feature_dim = spec.edge['feature'].shape[-1]
 
             if not self._update_edge_feature:
                 if (edge_feature_dim != node_feature_dim):
-                    warn(
+                    warnings.warn(
                         'Found edge feature dim to be incompatible with node feature dim. '
                         'Automatically adding a edge feature projection layer to match '
-                        'the dim of node features.'
+                        'the dim of node features.',
+                        stacklevel=2,
                     )
                     self._update_edge_feature = True 
 
@@ -603,19 +610,14 @@ class GIConv(GraphConv):
         else:
             self._update_edge_feature = False
 
-        has_overridden_update = self.__class__.update != GIConv.update 
-        if not has_overridden_update:
-            self._feedforward_intermediate_dense = self.get_dense(self.units)
-            self._feedforward_activation = self._activation
-            self._feedforward_output_dense = self.get_dense(self.units)
-
     def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         message = tensor.gather('feature', 'source')
         edge_feature = tensor.edge.get('feature')
         if self._update_edge_feature:
             edge_feature = self._edge_dense(edge_feature)
-        if self._has_edge_feature:
+        if self.has_edge_feature:
             message += edge_feature
+            message = keras.ops.relu(message)
         return tensor.update(
             {
                 'edge': {
@@ -639,20 +641,6 @@ class GIConv(GraphConv):
             }
         )
     
-    def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        node_feature = tensor.node['aggregate']
-        node_feature = self._feedforward_intermediate_dense(node_feature)
-        node_feature = self._feedforward_activation(node_feature)
-        node_feature = self._feedforward_output_dense(node_feature)
-        return tensor.update(
-            {
-                'node': {
-                    'feature': node_feature,
-                    'aggregate': None,
-                }
-            }
-        )
-    
     def get_config(self) -> dict:
         config = super().get_config()
         config.update({
@@ -701,15 +689,16 @@ class GAConv(GraphConv):
         activation: keras.layers.Activation | str | None = "relu",
         use_bias: bool = True,
         normalize: bool = False,
+        skip_connect: bool = True,
         update_edge_feature: bool = True,
-        attention_activation: keras.layers.Activation | str | None = "leaky_relu",
         **kwargs,
     ) -> None:
         super().__init__(
             units=units, 
             activation=activation,
-            use_bias=use_bias,
             normalize=normalize, 
+            use_bias=use_bias,
+            skip_connect=skip_connect,
             **kwargs
         )
         self._heads = heads
@@ -717,7 +706,6 @@ class GAConv(GraphConv):
             raise ValueError(f"units need to be divisible by heads.")
         self._head_units = self.units // self.heads 
         self._update_edge_feature = update_edge_feature
-        self._attention_activation = keras.activations.get(attention_activation)
 
     @property 
     def heads(self):
@@ -728,8 +716,7 @@ class GAConv(GraphConv):
         return self._head_units 
 
     def build(self, spec: tensors.GraphTensor.Spec) -> None:
-        self._has_edge_feature = 'feature' in spec.edge 
-        self._update_edge_feature = self._has_edge_feature and self._update_edge_feature
+        self._update_edge_feature = self.has_edge_feature and self._update_edge_feature
         if self._update_edge_feature:
             self._edge_dense = self.get_einsum_dense(
                 'ijh,jkh->ikh', (self.head_units, self.heads)
@@ -743,15 +730,6 @@ class GAConv(GraphConv):
         self._attention_dense = self.get_einsum_dense(
             'ijh,jkh->ikh', (1, self.heads)
         )
-        self._node_self_dense = self.get_einsum_dense(
-            'ij,jkh->ikh', (self.head_units, self.heads)
-        )
-
-        has_overridden_update = self.__class__.update != GAConv.update 
-        if not has_overridden_update:
-            self._feedforward_intermediate_dense = self.get_dense(self.units)
-            self._feedforward_activation = self._activation
-            self._feedforward_output_dense = self.get_dense(self.units)
 
     def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         attention_feature = keras.ops.concatenate(
@@ -761,7 +739,7 @@ class GAConv(GraphConv):
             ], 
             axis=-1
         )
-        if self._has_edge_feature:
+        if self.has_edge_feature:
             attention_feature = keras.ops.concatenate(
                 [
                     attention_feature, 
@@ -778,7 +756,7 @@ class GAConv(GraphConv):
             edge_feature = self._edge_dense(attention_feature)
             edge_feature = keras.ops.reshape(edge_feature, (-1, self.units))
         
-        attention_feature = self._attention_activation(attention_feature)
+        attention_feature = keras.ops.leaky_relu(attention_feature)
         attention_score = self._attention_dense(attention_feature)
         attention_score = ops.edge_softmax(
             score=attention_score, edge_target=tensor.edge['target']
@@ -797,7 +775,6 @@ class GAConv(GraphConv):
     
     def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         node_feature = tensor.aggregate('message', mode='sum')
-        node_feature += self._node_self_dense(tensor.node['feature'])
         node_feature = keras.ops.reshape(node_feature, (-1, self.units))
         return tensor.update(
             {
@@ -810,28 +787,123 @@ class GAConv(GraphConv):
             }
         )
     
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({
+            "heads": self._heads,
+            'update_edge_feature': self._update_edge_feature,
+        })
+        return config
+    
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class MPConv(GraphConv):
+
+    """Message passing neural network layer.
+
+    Also supports 3D molecular graphs.
+
+    >>> graph = molcraft.tensors.GraphTensor(
+    ...     context={
+    ...         'size': [2]
+    ...     },
+    ...     node={
+    ...         'feature': [[1.], [2.]]
+    ...     },
+    ...     edge={
+    ...         'source': [0, 1],
+    ...         'target': [1, 0],
+    ...     }
+    ... )
+    >>> conv = molcraft.layers.MPConv(units=4)
+    >>> conv(graph)
+        GraphTensor(
+            context={
+                'size': <tf.Tensor: shape=[1], dtype=int32>
+            },
+            node={
+                'feature': <tf.Tensor: shape=[2, 4], dtype=float32>
+            },
+            edge={
+                'source': <tf.Tensor: shape=[2], dtype=int32>,
+                'target': <tf.Tensor: shape=[2], dtype=int32>
+            }
+        )
+    """
+
+    def __init__(
+        self, 
+        units: int = 128, 
+        activation: keras.layers.Activation | str | None = 'relu', 
+        use_bias: bool = True,
+        normalize: bool = False,
+        skip_connect: bool = True,
+        **kwargs
+    ) -> None:
+        super().__init__(
+            units=units, 
+            activation=activation,
+            use_bias=use_bias,
+            normalize=normalize, 
+            skip_connect=skip_connect,
+            **kwargs
+        )
+
+    def build(self, spec: tensors.GraphTensor.Spec) -> None:
+        node_feature_dim = spec.node['feature'].shape[-1]
+        self.update_fn = keras.layers.GRUCell(self.units)
+        self._project_previous_node_feature = node_feature_dim != self.units
+        if self._project_previous_node_feature:
+            warnings.warn(
+                'Input node feature dim does not match updated node feature dim. '
+                'To make sure input node feature can be passed as `states` to the '
+                'GRU cell, it will automatically be projected prior to it.',
+                stacklevel=2
+            )
+            self._previous_node_dense = self.get_dense(self.units)
+
+    def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Aggregates messages.
+
+        This method may be overridden by subclass.
+
+        Arguments:
+            tensor:
+                A `GraphTensor` instance containing a message.
+        """
+        aggregate = tensor.aggregate('message', mode='mean')
+        return tensor.update(
+            {
+                'node': {
+                    'aggregate': aggregate, 
+                },
+                'edge': {
+                    'message': None,
+                }
+            }
+        )
+    
     def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        node_feature = tensor.node['aggregate']
-        node_feature = self._feedforward_intermediate_dense(node_feature)
-        node_feature = self._feedforward_activation(node_feature)
-        node_feature = self._feedforward_output_dense(node_feature)
+        previous = tensor.node['feature']
+        aggregate = tensor.node['aggregate']
+        if self._project_previous_node_feature:
+            previous = self._previous_node_dense(previous)
+        updated_node_feature, _ = self.update_fn(
+            inputs=aggregate, states=previous
+        )
         return tensor.update(
             {
                 'node': {
-                    'feature': node_feature,
+                    'feature': updated_node_feature,
                     'aggregate': None,
                 }
             }
         )
-    
+
     def get_config(self) -> dict:
         config = super().get_config()
-        config.update({
-            "heads": self._heads,
-            'update_edge_feature': self._update_edge_feature,
-            'attention_activation': keras.activations.serialize(self._attention_activation),
-        })
-        return config
+        config.update({})
+        return config 
     
 
 @keras.saving.register_keras_serializable(package='molcraft')
@@ -839,6 +911,8 @@ class GTConv(GraphConv):
 
     """Graph transformer layer.
 
+    Also supports 3D molecular graphs.
+
     >>> graph = molcraft.tensors.GraphTensor(
     ...     context={
     ...         'size': [2]
@@ -862,10 +936,9 @@ class GTConv(GraphConv):
             },
             edge={
                 'source': <tf.Tensor: shape=[2], dtype=int32>,
-                'target': <tf.Tensor: shape=[2], dtype=int32>
+                'target': <tf.Tensor: shape=[2], dtype=int32>,
             }
         )
-
     """
 
     def __init__(
@@ -875,14 +948,16 @@ class GTConv(GraphConv):
         activation: keras.layers.Activation | str | None = "relu",
         use_bias: bool = True,
         normalize: bool = False,
+        skip_connect: bool = True,
         attention_dropout: float = 0.0,
         **kwargs,
     ) -> None:
         super().__init__(
             units=units, 
             activation=activation,
-            use_bias=use_bias,
             normalize=normalize, 
+            use_bias=use_bias,
+            skip_connect=skip_connect,
             **kwargs
         )
         self._heads = heads
@@ -900,6 +975,8 @@ class GTConv(GraphConv):
         return self._head_units 
     
     def build(self, spec: tensors.GraphTensor.Spec) -> None:
+        """Builds the layer.
+        """
         self._query_dense = self.get_einsum_dense(
             'ij,jkh->ikh', (self.head_units, self.heads)
         )
@@ -912,29 +989,36 @@ class GTConv(GraphConv):
         self._output_dense = self.get_dense(self.units)
         self._softmax_dropout = keras.layers.Dropout(self._attention_dropout) 
 
-        self._add_bias = not 'bias' in spec.edge
-
-        if self._add_bias:
-            self._edge_bias = EdgeBias(biases=self.heads)
+        if self.has_edge_feature:
+            self._attention_bias_dense_1 = self.get_einsum_dense('ij,jkh->ikh', (1, self.heads))
 
-        has_overridden_update = self.__class__.update != GTConv.update 
-        if not has_overridden_update:
-            self._feedforward_intermediate_dense = self.get_dense(self.units)
-            self._feedforward_activation = self._activation
-            self._feedforward_output_dense = self.get_dense(self.units)
+        if self.has_node_coordinate:
+            node_feature_dim = spec.node['feature'].shape[-1]
+            num_kernels = self.units
+            self._gaussian_loc = self.add_weight(
+                shape=[num_kernels], initializer='zeros', dtype='float32', trainable=True
+            ) 
+            self._gaussian_scale = self.add_weight(
+                shape=[num_kernels], initializer='ones', dtype='float32', trainable=True
+            )
+            self._centrality_dense = self.get_dense(units=node_feature_dim)
+            self._attention_bias_dense_2 = self.get_einsum_dense('ij,jkh->ikh', (1, self.heads))
 
     def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        if self._add_bias:
-            edge_bias = self._edge_bias(tensor)
-            tensor = tensor.update(
-                {
-                    'edge': {
-                        'bias': edge_bias
-                    }
-                }
-            )
         node_feature = tensor.node['feature']
-        
+
+        if self.has_node_coordinate:
+            euclidean_distance = ops.euclidean_distance(
+                tensor.gather('coordinate', 'target'),
+                tensor.gather('coordinate', 'source'),
+                axis=-1
+            )
+            gaussian = ops.gaussian(
+                euclidean_distance, self._gaussian_loc, self._gaussian_scale
+            )
+            centrality = keras.ops.segment_sum(gaussian, tensor.edge['target'], tensor.num_nodes)
+            node_feature += self._centrality_dense(centrality)
+
         query = self._query_dense(node_feature)
         key = self._key_dense(node_feature)
         value = self._value_dense(node_feature)
@@ -946,23 +1030,45 @@ class GTConv(GraphConv):
         attention_score = keras.ops.sum(query * key, axis=1, keepdims=True)
         attention_score /= keras.ops.sqrt(float(self.head_units))
             
-        attention_score += keras.ops.expand_dims(tensor.edge['bias'], axis=1)
+        if self.has_edge_feature:
+            attention_score += self._attention_bias_dense_1(tensor.edge['feature'])
+
+        if self.has_node_coordinate:
+            attention_score += self._attention_bias_dense_2(gaussian)
+
         attention = ops.edge_softmax(attention_score, tensor.edge['target'])
         attention = self._softmax_dropout(attention)
-        message = ops.edge_weight(value, attention)
 
+        if self.has_node_coordinate:
+            displacement = ops.displacement(
+                tensor.gather('coordinate', 'target'),
+                tensor.gather('coordinate', 'source'),
+                normalize=True
+            )
+            attention *= keras.ops.expand_dims(displacement, axis=-1)
+            attention = keras.ops.expand_dims(attention, axis=2)
+            value = keras.ops.expand_dims(value, axis=1)
+
+        message = ops.edge_weight(value, attention)
+        
         return tensor.update(
             {
                 'edge': {
-                    'message': message
+                    'message': message,
                 },
             }
         )
 
     def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         node_feature = tensor.aggregate('message', mode='sum')
-        node_feature = keras.ops.reshape(node_feature, (-1, self.units))
+        if self.has_node_coordinate:
+            shape = (tensor.num_nodes, -1, self.units)
+        else:
+            shape = (tensor.num_nodes, self.units)
+        node_feature = keras.ops.reshape(node_feature, shape)
         node_feature = self._output_dense(node_feature)
+        if self.has_node_coordinate:
+            node_feature = keras.ops.sum(node_feature, axis=1)
         return tensor.update(
             {
                 'node': {
@@ -973,20 +1079,6 @@ class GTConv(GraphConv):
                 }
             }
         )
-
-    def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        node_feature = tensor.node['aggregate']
-        node_feature = self._feedforward_intermediate_dense(node_feature)
-        node_feature = self._feedforward_activation(node_feature)
-        node_feature = self._feedforward_output_dense(node_feature)
-        return tensor.update(
-            {
-                'node': {
-                    'feature': node_feature,
-                    'aggregate': None,
-                },
-            }
-        )
     
     def get_config(self) -> dict:
         config = super().get_config()
@@ -998,17 +1090,49 @@ class GTConv(GraphConv):
     
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class MPConv(GraphConv):
+class EGConv(GraphConv):
 
-    """Message passing neural network layer.
+    """Equivariant graph neural network layer 3D.
+
+    Only supports 3D molecular graphs.
+
+    >>> graph = molcraft.tensors.GraphTensor(
+    ...     context={
+    ...         'size': [2]
+    ...     },
+    ...     node={
+    ...         'feature': [[1.], [2.]],
+    ...         'coordinate': [[0.1, -0.1, 0.5], [1.2, -0.5, 2.1]],
+    ...     },
+    ...     edge={
+    ...         'source': [0, 1],
+    ...         'target': [1, 0],
+    ...     }
+    ... )
+    >>> conv = molcraft.layers.EGConv(units=4)
+    >>> conv(graph)
+        GraphTensor(
+            context={
+                'size': <tf.Tensor: shape=[1], dtype=int32>
+            },
+            node={
+                'feature': <tf.Tensor: shape=[2, 4], dtype=float32>,
+                'coordinate': <tf.Tensor: shape=[2, 3], dtype=float32>
+            },
+            edge={
+                'source': <tf.Tensor: shape=[2], dtype=int32>,
+                'target': <tf.Tensor: shape=[2], dtype=int32>
+            }
+        )
     """
 
     def __init__(
         self, 
         units: int = 128, 
-        activation: keras.layers.Activation | str | None = None, 
+        activation: keras.layers.Activation | str | None = 'silu', 
         use_bias: bool = True,
         normalize: bool = False,
+        skip_connect: bool = True,
         **kwargs
     ) -> None:
         super().__init__(
@@ -1016,262 +1140,30 @@ class MPConv(GraphConv):
             activation=activation,
             use_bias=use_bias,
             normalize=normalize, 
+            skip_connect=skip_connect,
             **kwargs
         )
 
     def build(self, spec: tensors.GraphTensor.Spec) -> None:
-        node_feature_dim = spec.node['feature'].shape[-1]
-        self.message_fn = self.get_dense(self.units, activation=self._activation)
-        self.update_fn = keras.layers.GRUCell(self.units)
-        self._has_edge_feature = 'feature' in spec.edge
-        self.project_input_node_feature = node_feature_dim != self.units
-        if self.project_input_node_feature:
-            warn(
-                'Input node feature dim does not match updated node feature dim. '
-                'To make sure input node feature can be passed as `states` to the '
-                'GRU cell, it will automatically be projected prior to it.'
-            )
-            self._previous_node_dense = self.get_dense(
-                self.units, activation=self._activation
+        if not self.has_node_coordinate:
+            raise ValueError(
+                'Could not find `coordinate`s in node, '
+                'which is required for Conv3D layers.'
             )
-
-    def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        feature = keras.ops.concatenate(
-            [
-                tensor.gather('feature', 'source'),
-                tensor.gather('feature', 'target'),
-            ], 
-            axis=-1
-        )
-        if self._has_edge_feature:
-            feature = keras.ops.concatenate(
-                [
-                    feature,
-                    tensor.edge['feature']
-                ], 
-                axis=-1
-            )
-        message = self.message_fn(feature)
-        return tensor.update(
-            {
-                'edge': {
-                    'message': message,
-                }
-            }
-        )
-
-    def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        aggregate = tensor.aggregate('message', mode='mean')
-        feature = tensor.node['feature']
-        if self.project_input_node_feature:
-            feature = self._previous_node_dense(feature)
-        return tensor.update(
-            {
-                'node': {
-                    'aggregate': aggregate,
-                    'feature': feature,
-                }
-            }
-        )
-
-    def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        updated_node_feature, _ = self.update_fn(
-            inputs=tensor.node['aggregate'],
-            states=tensor.node['feature']
-        )
-        return tensor.update(
-            {
-                'node': {
-                    'feature': updated_node_feature,
-                    'aggregate': None,
-                }
-            }
-        )
-
-    def get_config(self) -> dict:
-        config = super().get_config()
-        config.update({})
-        return config 
-    
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class GTConv3D(GTConv):
-
-    """Graph transformer layer 3D.
-    """
-
-    def build(self, spec: tensors.GraphTensor.Spec) -> None:
-        """Builds the layer.
-        """
-        super().build(spec)
-        if self._add_bias:
-            node_feature_dim = spec.node['feature'].shape[-1]
-            kernels = self.units
-            self._gaussian_basis = GaussianDistance(kernels)
-            self._centrality_dense = self.get_dense(units=node_feature_dim)
-            self._gaussian_edge_bias = self.get_dense(self.heads)
-
-    def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        node_feature = tensor.node['feature']
-
-        if self._add_bias:
-            gaussian = self._gaussian_basis(tensor)
-            centrality = keras.ops.segment_sum(
-                gaussian, tensor.edge['target'], tensor.num_nodes
-            )
-            node_feature += self._centrality_dense(centrality)
-
-            edge_bias = self._edge_bias(tensor) + self._gaussian_edge_bias(gaussian)
-            tensor = tensor.update({'edge': {'bias': edge_bias}})
-            
-        query = self._query_dense(node_feature)
-        key = self._key_dense(node_feature)
-        value = self._value_dense(node_feature)
-
-        query = ops.gather(query, tensor.edge['source'])
-        key = ops.gather(key, tensor.edge['target'])
-        value = ops.gather(value, tensor.edge['source'])
-
-        attention_score = keras.ops.sum(query * key, axis=1, keepdims=True)
-        attention_score /= keras.ops.sqrt(float(self.head_units))
-            
-        attention_score += keras.ops.expand_dims(tensor.edge['bias'], axis=1)
-        
-        attention = ops.edge_softmax(attention_score, tensor.edge['target'])
-        attention = self._softmax_dropout(attention)
-
-        distance = keras.ops.subtract(
-            tensor.gather('coordinate', 'source'),
-            tensor.gather('coordinate', 'target')
-        )
-        euclidean_distance = ops.euclidean_distance(
-            tensor.gather('coordinate', 'source'),
-            tensor.gather('coordinate', 'target'),
-            axis=-1
-        )
-        distance /= euclidean_distance
-
-        attention *= keras.ops.expand_dims(distance, axis=-1)
-        attention = keras.ops.expand_dims(attention, axis=2)
-        value = keras.ops.expand_dims(value, axis=1)
-
-        message = ops.edge_weight(value, attention)
-        
-        return tensor.update(
-            {
-                'edge': {
-                    'message': message,
-                },
-            }
-        )
-
-    def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        node_feature = tensor.aggregate('message', mode='sum')
-        node_feature = keras.ops.reshape(
-            node_feature, (tensor.num_nodes, -1, self.units)
-        )
-        node_feature = self._output_dense(node_feature)
-        node_feature = keras.ops.sum(node_feature, axis=1)
-        return tensor.update(
-            {
-                'node': {
-                    'aggregate': node_feature,
-                },
-                'edge': {
-                    'message': None,
-                }
-            }
-        )
-    
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class MPConv3D(MPConv):
-
-    """Message passing neural network layer 3D.
-    """
-    
-    def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        euclidean_distance = ops.euclidean_distance(
-            tensor.gather('coordinate', 'target'),
-            tensor.gather('coordinate', 'source'),
-            axis=-1
-        )
-        feature = keras.ops.concatenate(
-            [
-                tensor.gather('feature', 'source'),
-                tensor.gather('feature', 'target'),
-                euclidean_distance,
-            ], 
-            axis=-1
-        )
-        if self._has_edge_feature:
-            feature = keras.ops.concatenate(
-                [
-                    feature,
-                    tensor.edge['feature']
-                ], 
-                axis=-1
-            )
-        message = self.message_fn(feature)
-        return tensor.update(
-            {
-                'edge': {
-                    'message': message,
-                }
-            }
-        )
-    
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class EGConv3D(GraphConv):
-
-    """Equivariant graph neural network layer 3D.
-    """
-
-    def __init__(
-        self, 
-        units: int = 128, 
-        activation: keras.layers.Activation | str | None = 'silu', 
-        use_bias: bool = True,
-        normalize: bool = False,
-        **kwargs
-    ) -> None:
-        super().__init__(
-            units=units, 
-            activation=activation,
-            use_bias=use_bias,
-            normalize=normalize, 
-            **kwargs
-        )
-
-    def build(self, spec: tensors.GraphTensor.Spec) -> None:
-        if 'coordinate' not in spec.node:
-            raise ValueError(
-                'Could not find `coordinate`s in node, '
-                'which is required for Conv3D layers.'
-            )
-        self._has_edge_feature = 'feature' in spec.edge
-        self._message_feedforward_intermediate = self.get_dense(
-            self.units, activation=self._activation
-        )
-        self._message_feedforward_final = self.get_dense(
-            self.units, activation=self._activation
-        )
+        self._message_feedforward_intermediate = self.get_dense(
+            self.units, activation=self.activation
+        )
+        self._message_feedforward_final = self.get_dense(
+            self.units, activation=self.activation
+        )
 
         self._coord_feedforward_intermediate = self.get_dense(
-            self.units, activation=self._activation
+            self.units, activation=self.activation
         ) 
         self._coord_feedforward_final = self.get_dense(
             1, use_bias=False, activation='tanh'
         )
 
-        has_overridden_update = self.__class__.update != EGConv3D.update 
-        if not has_overridden_update:
-            self._feedforward_intermediate = self.get_dense(
-                self.units, activation=self._activation
-            )
-            self._feedforward_output = self.get_dense(self.units)
-
     def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         relative_node_coordinate = keras.ops.subtract(
             tensor.gather('coordinate', 'target'), 
@@ -1300,7 +1192,7 @@ class EGConv3D(GraphConv):
             ], 
             axis=-1
         )
-        if self._has_edge_feature:
+        if self.has_edge_feature:
             feature = keras.ops.concatenate(
                 [
                     feature,
@@ -1339,7 +1231,7 @@ class EGConv3D(GraphConv):
         # graph to graph). Therefore, a mean mean aggregation is performed 
         # instead:
         aggregate = tensor.aggregate('message', mode='mean')
-
+        aggregate = keras.ops.concatenate([aggregate, tensor.node['feature']], axis=-1)
         # Simply added to silence warning ('no gradients for variables ...')
         aggregate += (0.0 * keras.ops.sum(coordinate))
 
@@ -1355,26 +1247,6 @@ class EGConv3D(GraphConv):
                 }
             }
         ) 
-
-    def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        feature = keras.ops.concatenate(
-            [
-                tensor.node['aggregate'], 
-                tensor.node['feature']
-            ], 
-            axis=-1
-        )
-        updated_node_feature = self._feedforward_output(
-            self._feedforward_intermediate(feature)
-        )  
-        return tensor.update(
-            {
-                'node': {
-                    'feature': updated_node_feature, 
-                    'aggregate': None,
-                },
-            }
-        )
     
     def get_config(self) -> dict:
         config = super().get_config()
@@ -1417,146 +1289,6 @@ class Readout(GraphLayer):
         config['mode'] = self.mode 
         return config 
     
-    
-@keras.saving.register_keras_serializable(package='molcraft')
-class GraphNetwork(GraphLayer):
-
-    """Graph neural network.
-
-    Sequentially calls graph layers (`GraphLayer`) and concatenates its output. 
-
-    Arguments:
-        layers (list):
-            A list of graph layers.
-    """
-
-    def __init__(self, layers: list[GraphLayer], **kwargs) -> None:
-        super().__init__(**kwargs)
-        self.layers = layers
-        self._update_edge_feature = False
-
-    def build(self, spec: tensors.GraphTensor.Spec) -> None:
-        units = self.layers[0].units 
-        node_feature_dim = spec.node['feature'].shape[-1]
-        self._update_node_feature = node_feature_dim != units 
-        if self._update_node_feature:
-            warn(
-                'Node feature dim does not match `units` of the first layer. '
-                'Automatically adding a node projection layer to match `units`.'
-            )
-            self._node_dense = self.get_dense(units)
-        self._has_edge_feature = 'feature' in spec.edge 
-        if self._has_edge_feature:
-            edge_feature_dim = spec.edge['feature'].shape[-1]
-            self._update_edge_feature = edge_feature_dim != units
-            if self._update_edge_feature:
-                warn(
-                    'Edge feature dim does not match `units` of the first layer. '
-                    'Automatically adding a edge projection layer to match `units`.'
-                )
-                self._edge_dense = self.get_dense(units)
-
-    def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        x = tensors.to_dict(tensor)
-        if self._update_node_feature:
-            x['node']['feature'] = self._node_dense(tensor.node['feature'])
-        if self._has_edge_feature and self._update_edge_feature:
-            x['edge']['feature'] = self._edge_dense(tensor.edge['feature'])
-        outputs = [x['node']['feature']]
-        for layer in self.layers:
-            x = layer(x)
-            outputs.append(x['node']['feature'])
-        return tensor.update(
-            {
-                'node': {
-                    'feature': keras.ops.concatenate(outputs, axis=-1)
-                }    
-            }
-        )
-    
-    def tape_propagate(
-        self,
-        tensor: tensors.GraphTensor,
-        tape: tf.GradientTape,
-        training: bool | None = None,
-    ) -> tuple[tensors.GraphTensor, list[tf.Tensor]]:
-        """Performs the propagation with a `GradientTape`.
-
-        Performs the same forward pass as `propagate` but with a `GradientTape`
-        watching intermediate node features.
-
-        Arguments:
-            tensor (tensors.GraphTensor):
-                The graph input.
-        """
-        if isinstance(tensor, tensors.GraphTensor):
-            x = tensors.to_dict(tensor)
-        else:
-            x = tensor
-        if self._update_node_feature:
-            x['node']['feature'] = self._node_dense(tensor.node['feature'])
-        if self._update_edge_feature:
-            x['edge']['feature'] = self._edge_dense(tensor.edge['feature'])
-        tape.watch(x['node']['feature'])
-        outputs = [x['node']['feature']]
-        for layer in self.layers:
-            x = layer(x, training=training)
-            tape.watch(x['node']['feature'])
-            outputs.append(x['node']['feature'])
-
-        tensor = tensor.update(
-            {
-                'node': {
-                    'feature': keras.ops.concatenate(outputs, axis=-1)
-                }
-            }
-        )
-        return tensor, outputs
-    
-    def get_config(self) -> dict:
-        config = super().get_config()
-        config.update(
-            {
-                'layers': [
-                    keras.layers.serialize(layer) for layer in self.layers
-                ]
-            }
-        )
-        return config
-
-    @classmethod
-    def from_config(cls, config: dict) -> 'GraphNetwork':
-        config['layers'] = [
-            keras.layers.deserialize(layer) for layer in config['layers']
-        ]
-        return super().from_config(config)
-    
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class Extraction(GraphLayer):
-
-    def __init__(
-        self, 
-        field: str, 
-        inner_field: str | None = None, 
-        **kwargs
-    ) -> None:
-        super().__init__(**kwargs)
-        self.field = field 
-        self.inner_field = inner_field
-
-    def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        data = dict(getattr(tensor, self.field))
-        if not self.inner_field:
-            return data 
-        return data[self.inner_field]
-     
-    def get_config(self):
-        config = super().get_config()
-        config['field'] = self.field
-        config['inner_field'] = self.inner_field 
-        return config 
-    
 
 @keras.saving.register_keras_serializable(package='molcraft')
 class NodeEmbedding(GraphLayer):
@@ -1571,17 +1303,12 @@ class NodeEmbedding(GraphLayer):
         dim: int = None, 
         normalize: bool = False,
         embed_context: bool = False,
-        allow_reconstruction: bool = False,
-        allow_masking: bool = False, 
         **kwargs
     ) -> None:
         super().__init__(**kwargs)
         self.dim = dim
         self._normalize = normalize
         self._embed_context = embed_context
-        self._masking_rate = None
-        self._allow_masking = allow_masking
-        self._allow_reconstruction = allow_reconstruction
 
     def build(self, spec: tensors.GraphTensor.Spec) -> None:
         feature_dim = spec.node['feature'].shape[-1]
@@ -1595,47 +1322,31 @@ class NodeEmbedding(GraphLayer):
             self._embed_context = False 
         if self._has_super and not self._embed_context:
             self._super_feature = self.get_weight(shape=[self.dim], name='super_node_feature')
-        if self._allow_masking:
-            self._mask_feature = self.get_weight(shape=[self.dim], name='mask_node_feature')
         if self._embed_context:
             self._context_dense = self.get_dense(self.dim)
         
-        if self._normalize:
-            if str(self._normalize).lower().startswith('batch'):
-                self._norm = keras.layers.BatchNormalization(
-                    name='output_batch_norm'
-                )
-            else:
-                self._norm = keras.layers.LayerNormalization(
-                    name='output_layer_norm'
-                )
+        if not self._normalize:
+            self._norm = keras.layers.Identity()
+        elif str(self._normalize).lower().startswith('layer'):
+            self._norm = keras.layers.LayerNormalization()
+        else:
+            self._norm = keras.layers.BatchNormalization()
 
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         feature = self._node_dense(tensor.node['feature'])
 
-        if self._has_super:
-            super_feature = (0 if self._embed_context else self._super_feature)
+        if self._has_super and not self._embed_context:
             super_mask = keras.ops.expand_dims(tensor.node['super'], 1)
-            feature = keras.ops.where(super_mask, super_feature, feature)
+            feature = keras.ops.where(super_mask, self._super_feature, feature)
 
         if self._embed_context:
             context_feature = self._context_dense(tensor.context['feature'])
             feature = ops.scatter_update(feature, tensor.node['super'], context_feature)
             tensor = tensor.update({'context': {'feature': None}})
 
-        apply_mask = (self._allow_masking and 'mask' in tensor.node)
-        if apply_mask:
-            mask = keras.ops.expand_dims(tensor.node['mask'], -1)
-            feature = keras.ops.where(mask, self._mask_feature, feature)
-        elif self._allow_masking:
-            feature = feature + (self._mask_feature * 0.0)
+        feature = self._norm(feature)
 
-        if self._normalize:
-            feature = self._norm(feature)
-
-        if not self._allow_reconstruction:
-            return tensor.update({'node': {'feature': feature}})
-        return tensor.update({'node': {'feature': feature, 'target_feature': feature}})
+        return tensor.update({'node': {'feature': feature}})
 
     def get_config(self) -> dict:
         config = super().get_config()
@@ -1643,8 +1354,6 @@ class NodeEmbedding(GraphLayer):
             'dim': self.dim,
             'normalize': self._normalize,
             'embed_context': self._embed_context,
-            'allow_masking': self._allow_masking,
-            'allow_reconstruction': self._allow_reconstruction,
         })
         return config
     
@@ -1661,39 +1370,30 @@ class EdgeEmbedding(GraphLayer):
         self, 
         dim: int = None, 
         normalize: bool = False,
-        allow_masking: bool = True, 
         **kwargs
     ) -> None:
         super().__init__(**kwargs)
         self.dim = dim
         self._normalize = normalize
-        self._masking_rate = None
-        self._allow_masking = allow_masking
 
     def build(self, spec: tensors.GraphTensor.Spec) -> None:
         feature_dim = spec.edge['feature'].shape[-1]
         if not self.dim:
             self.dim = feature_dim
-        self._edge_dense = self.get_dense(self.dim)
+        self._edge_dense = self.get_dense(self.dim) 
+
+        self._self_loop_feature = self.get_weight(shape=[self.dim], name='self_loop_edge_feature')
 
         self._has_super = 'super' in spec.edge
-        self._has_self_loop = 'self_loop' in spec.edge
         if self._has_super:
             self._super_feature = self.get_weight(shape=[self.dim], name='super_edge_feature')
-        if self._has_self_loop:
-            self._self_loop_feature = self.get_weight(shape=[self.dim], name='self_loop_edge_feature')
-        if self._allow_masking:
-            self._mask_feature = self.get_weight(shape=[self.dim], name='mask_edge_feature')
-
-        if self._normalize:
-            if str(self._normalize).lower().startswith('batch'):
-                self._norm = keras.layers.BatchNormalization(
-                    name='output_batch_norm'
-                )
-            else:
-                self._norm = keras.layers.LayerNormalization(
-                    name='output_layer_norm'
-                )
+        
+        if not self._normalize:
+            self._norm = keras.layers.Identity()
+        elif str(self._normalize).lower().startswith('layer'):
+            self._norm = keras.layers.LayerNormalization()
+        else:
+            self._norm = keras.layers.BatchNormalization()
 
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         feature = self._edge_dense(tensor.edge['feature'])
@@ -1702,246 +1402,136 @@ class EdgeEmbedding(GraphLayer):
             super_mask = keras.ops.expand_dims(tensor.edge['super'], 1)
             feature = keras.ops.where(super_mask, self._super_feature, feature)
 
-        if self._has_self_loop:
-            self_loop_mask = keras.ops.expand_dims(tensor.edge['self_loop'], 1)
-            feature = keras.ops.where(self_loop_mask, self._self_loop_feature, feature)
-            
-        if (
-            self._allow_masking and 
-            self._masking_rate is not None and 
-            self._masking_rate > 0
-        ):
-            random = keras.random.uniform(shape=[tensor.num_edges])
-            mask = random <= self._masking_rate
-            if self._has_super:
-                mask = keras.ops.logical_and(
-                    mask, keras.ops.logical_not(tensor.edge['super'])
-                )
-            mask = keras.ops.expand_dims(mask, -1)
-            feature = keras.ops.where(mask, self._mask_feature, feature)
-        elif self._allow_masking:
-            # Simply added to silence warning ('no gradients for variables ...')
-            feature += (0.0 * self._mask_feature)
+        self_loop_mask = keras.ops.expand_dims(tensor.edge['source'] == tensor.edge['target'], 1)
+        feature = keras.ops.where(self_loop_mask, self._self_loop_feature, feature)
 
-        if self._normalize:
-            feature = self._norm(feature)
+        feature = self._norm(feature)
 
-        return tensor.update({'edge': {'feature': feature, 'embedding': feature}})
-
-    @property 
-    def masking_rate(self):
-        return self._masking_rate 
-    
-    @masking_rate.setter
-    def masking_rate(self, rate: float):
-        if not self._allow_masking and rate is not None:
-            raise ValueError(
-                f'Cannot set `masking_rate` for layer {self} '
-                'as `allow_masking` was set to `False`.'
-            )
-        self._masking_rate = float(rate)
+        return tensor.update({'edge': {'feature': feature}})
 
     def get_config(self) -> dict:
         config = super().get_config()
         config.update({
             'dim': self.dim,
             'normalize': self._normalize,
-            'allow_masking': self._allow_masking
         })
         return config
     
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class Projection(GraphLayer):
-    """Base graph projection layer.
-    """
-    def __init__(
-        self, 
-        units: int = None, 
-        activation: str | keras.layers.Activation | None = None, 
-        use_bias: bool = True,
-        field: str = 'node',
-        **kwargs
-    ) -> None:
-        super().__init__(use_bias=use_bias, **kwargs)
-        self.units = units
-        self._activation = keras.activations.get(activation)
-        self.field = field 
-
-    def build(self, spec: tensors.GraphTensor.Spec) -> None:
-        data = getattr(spec, self.field, None)
-        if data is None:
-            raise ValueError('Could not access field {self.field!r}.')
-        feature_dim = data['feature'].shape[-1]
-        if not self.units:
-            self.units = feature_dim
-        self._dense = self.get_dense(self.units)
-
-    def propagate(self, tensor: tensors.GraphTensor):
-        feature = getattr(tensor, self.field)['feature']
-        feature = self._dense(feature)
-        feature = self._activation(feature)
-        return tensor.update(
-            {
-                self.field: {
-                    'feature': feature
-                }
-            }
-        ) 
-
-    def get_config(self) -> dict:
-        config = super().get_config()
-        config.update({
-            'units': self.units,
-            'activation': keras.activations.serialize(self._activation),
-            'field': self.field,
-        })
-        return config
-    
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class ContextProjection(Projection):
-    """Context projection layer.
-    """
-    def __init__(self, units: int = None, activation: str = None, **kwargs):
-        kwargs['field'] = 'context'
-        super().__init__(units=units, activation=activation, **kwargs)
-
+class GraphNetwork(GraphLayer):
 
-@keras.saving.register_keras_serializable(package='molcraft')
-class NodeProjection(Projection):
-    """Node projection layer.
-    """
-    def __init__(self, units: int = None, activation: str = None, **kwargs):
-        kwargs['field'] = 'node'
-        super().__init__(units=units, activation=activation, **kwargs)
+    """Graph neural network.
 
+    Sequentially calls graph layers (`GraphLayer`) and concatenates its output. 
 
-@keras.saving.register_keras_serializable(package='molcraft')
-class EdgeProjection(Projection):
-    """Edge projection layer.
+    Arguments:
+        layers (list):
+            A list of graph layers.
     """
-    def __init__(self, units: int = None, activation: str = None, **kwargs):
-        kwargs['field'] = 'edge'
-        super().__init__(units=units, activation=activation, **kwargs)
-
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class Reconstruction(GraphLayer):
-
-    def __init__(
-        self, 
-        loss: keras.losses.Loss | str = 'mse',
-        loss_weight: float = 0.5, 
-        **kwargs
-    ):
-        super().__init__(**kwargs)
-        self._loss_fn = keras.losses.get(loss)
-        self._loss_weight = loss_weight
-    
-    def build(self, spec: tensors.GraphTensor.Spec) -> None:
-         has_target_node_feature = 'target_feature' in spec.node
-         if not has_target_node_feature:
-             raise ValueError(
-                'Could not find `target_feature` in `spec.node`. '
-                'Add a `target_feature` via `NodeEmbedding` by setting '
-                '`allow_reconstruction` to `True`.'
-            )
-         output_dim = spec.node['target_feature'].shape[-1]
-         self._dense = self.get_dense(output_dim)
-
-    def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        target_node_feature = tensor.node['target_feature']
-        transformed_node_feature = tensor.node['feature']
-
-        reconstructed_node_feature = self._dense(
-            transformed_node_feature
-        )
-
-        loss = self._loss_fn(
-            target_node_feature, reconstructed_node_feature
-        ) 
-        self.add_loss(keras.ops.sum(loss) * self._loss_weight)
-        return tensor.update({'node': {'feature': transformed_node_feature}})
-    
-    def get_config(self):
-        config = super().get_config()
-        config['loss'] = keras.losses.serialize(self._loss_fn)
-        config['loss_weight'] = self._loss_weight 
-        return config
-    
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class EdgeBias(GraphLayer):
 
-    def __init__(self, biases: int, **kwargs):
+    def __init__(self, layers: list[GraphLayer], **kwargs) -> None:
         super().__init__(**kwargs)
-        self.biases = biases
+        self.layers = layers
+        self._update_edge_feature = False
 
     def build(self, spec: tensors.GraphTensor.Spec) -> None:
-        self._has_edge_length = 'length' in spec.edge
-        self._has_edge_feature = 'feature' in spec.edge
-        if self._has_edge_feature:
-            self._edge_feature_dense = self.get_dense(self.biases)
-        if self._has_edge_length:
-            self._edge_length_dense = self.get_dense(
-                self.biases, kernel_initializer='zeros'
+        units = self.layers[0].units 
+        node_feature_dim = spec.node['feature'].shape[-1]
+        self._update_node_feature = node_feature_dim != units 
+        if self._update_node_feature:
+            warnings.warn(
+                'Node feature dim does not match `units` of the first layer. '
+                'Automatically adding a node projection layer to match `units`.',
+                stacklevel=2
             )
+            self._node_dense = self.get_dense(units)
+        self._has_edge_feature = 'feature' in spec.edge 
+        if self._has_edge_feature:
+            edge_feature_dim = spec.edge['feature'].shape[-1]
+            self._update_edge_feature = edge_feature_dim != units
+            if self._update_edge_feature:
+                warnings.warn(
+                    'Edge feature dim does not match `units` of the first layer. '
+                    'Automatically adding a edge projection layer to match `units`.',
+                    stacklevel=2
+                )
+                self._edge_dense = self.get_dense(units)
 
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        bias = keras.ops.zeros(
-            shape=(tensor.num_edges, self.biases), 
-            dtype=tensor.node['feature'].dtype
+        x = tensors.to_dict(tensor)
+        if self._update_node_feature:
+            x['node']['feature'] = self._node_dense(tensor.node['feature'])
+        if self._has_edge_feature and self._update_edge_feature:
+            x['edge']['feature'] = self._edge_dense(tensor.edge['feature'])
+        outputs = [x['node']['feature']]
+        for layer in self.layers:
+            x = layer(x)
+            outputs.append(x['node']['feature'])
+        return tensor.update(
+            {
+                'node': {
+                    'feature': keras.ops.concatenate(outputs, axis=-1)
+                }    
+            }
         )
-        if self._has_edge_feature:
-            bias += self._edge_feature_dense(tensor.edge['feature'])
-        if self._has_edge_length:
-            bias += self._edge_length_dense(tensor.edge['length'])
-        return bias
-
-    def get_config(self) -> dict:
-        config = super().get_config()
-        config.update({'biases': self.biases})
-        return config
     
+    def tape_propagate(
+        self,
+        tensor: tensors.GraphTensor,
+        tape: tf.GradientTape,
+        training: bool | None = None,
+    ) -> tuple[tensors.GraphTensor, list[tf.Tensor]]:
+        """Performs the propagation with a `GradientTape`.
 
-@keras.saving.register_keras_serializable(package='molcraft')
-class GaussianDistance(GraphLayer):
-
-    def __init__(self, kernels: int, **kwargs):
-        super().__init__(**kwargs)
-        self.kernels = kernels 
+        Performs the same forward pass as `propagate` but with a `GradientTape`
+        watching intermediate node features.
 
-    def build(self, spec: tensors.GraphTensor.Spec) -> None:
-        self._loc = self.add_weight(
-            shape=[self.kernels],
-            initializer='zeros',
-            dtype='float32',
-            trainable=True
-        ) 
-        self._scale = self.add_weight(
-            shape=[self.kernels],
-            initializer='ones',
-            dtype='float32',
-            trainable=True
-        )
+        Arguments:
+            tensor (tensors.GraphTensor):
+                The graph input.
+        """
+        if isinstance(tensor, tensors.GraphTensor):
+            x = tensors.to_dict(tensor)
+        else:
+            x = tensor
+        if self._update_node_feature:
+            x['node']['feature'] = self._node_dense(tensor.node['feature'])
+        if self._update_edge_feature:
+            x['edge']['feature'] = self._edge_dense(tensor.edge['feature'])
+        tape.watch(x['node']['feature'])
+        outputs = [x['node']['feature']]
+        for layer in self.layers:
+            x = layer(x, training=training)
+            tape.watch(x['node']['feature'])
+            outputs.append(x['node']['feature'])
 
-    def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        euclidean_distance = ops.euclidean_distance(
-            tensor.gather('coordinate', 'source'),
-            tensor.gather('coordinate', 'target'),
-            axis=-1
-        )
-        return ops.gaussian(
-            euclidean_distance, self._loc, self._scale
+        tensor = tensor.update(
+            {
+                'node': {
+                    'feature': keras.ops.concatenate(outputs, axis=-1)
+                }
+            }
         )
-
+        return tensor, outputs
+    
     def get_config(self) -> dict:
         config = super().get_config()
-        config.update({
-            'kernels': self.kernels,
-        })
+        config.update(
+            {
+                'layers': [
+                    keras.layers.serialize(layer) for layer in self.layers
+                ]
+            }
+        )
         return config
+
+    @classmethod
+    def from_config(cls, config: dict) -> 'GraphNetwork':
+        config['layers'] = [
+            keras.layers.deserialize(layer) for layer in config['layers']
+        ]
+        return super().from_config(config)
     
 
 @keras.saving.register_keras_serializable(package='molcraft')
@@ -1992,7 +1582,7 @@ class GaussianParams(keras.layers.Dense):
         config['units'] = None
         config['activation'] = keras.activations.serialize(self.loc_activation)
         return config
-
+    
 
 def Input(spec: tensors.GraphTensor.Spec) -> dict:
     """Used to specify inputs to model.
@@ -2047,13 +1637,6 @@ def Input(spec: tensors.GraphTensor.Spec) -> dict:
     return inputs
 
 
-def warn(message: str) -> None:
-    warnings.warn(
-        message=message,
-        category=UserWarning,
-        stacklevel=1
-    )
-
 def _serialize_spec(spec: tensors.GraphTensor.Spec) -> dict:
     serialized_spec = {}
     for outer_field, data in spec.__dict__.items():
@@ -2095,5 +1678,3 @@ def _spec_from_inputs(inputs):
 
 
 GraphTransformer = GTConv
-GraphTransformer3D = GTConv3D
-