molcraft 0.1.0a1__py3-none-any.whl → 0.1.0a3__py3-none-any.whl

molcraft/layers.py

@@ -60,25 +60,20 @@ class GraphLayer(keras.layers.Layer):
         May use built-in methods such as `get_weight`, `get_dense` and `get_einsum_dense`.
 
         Optionally implemented by subclass. If implemented, it is recommended to
-        build the sub-layers via `build([None, input_dim])`. If sub-layers are not 
-        built, symbolic input will be passed through the layer to build it.
+        If sub-layers are built (via `build` or `build_from_spec`), set `built`
+        to True. If not, symbolic input will be passed through the layer to build them.
 
         Args:
             spec:
-                A `GraphTensor.Spec` instance, corresponding to the input `GraphTensor` 
-                of the `propagate` method.
+                A `GraphTensor.Spec` instance, corresponding to the `GraphTensor` 
+                passed to `propagate`.
         """ 
-
+    
     def build(self, spec: tensors.GraphTensor.Spec) -> None:
             
         self._custom_build_config = {'spec': _serialize_spec(spec)}
 
-        invoke_build_from_spec = (
-            GraphLayer.build_from_spec != self.__class__.build_from_spec
-        )
-        if invoke_build_from_spec:
-            self.build_from_spec(spec)
-            self.built = True
+        self.build_from_spec(spec)
 
         if not self.built:
             # Automatically build layer or model by calling it on symbolic inputs
@@ -206,12 +201,66 @@ class GraphLayer(keras.layers.Layer):
 class GraphConv(GraphLayer):
 
     """Base graph neural network layer.
+
+    For normalization and skip connection to work, the `GraphConv` subclass 
+    requires the (node feature) output of `aggregate` and `update` to have a 
+    dimension of `self.units`, respectively. 
+
+    Args:
+        units:
+            The number of units.
+        normalize:
+            Whether `LayerNormalization` should be applied to the (node feature) output 
+            of the `aggregate` step. While normalization is recommended, it is not used 
+            by default.
+        skip_connection:
+            Whether (node feature) input should be added to the (node feature) output. 
+            If (node feature) input dim is not equal to `units`, a projection layer will 
+            automatically project the residual before adding it to the output. While skip 
+            connection is recommended, it is not used by default. 
+        kwargs:
+            See arguments of `GraphLayer`.
     """
         
-    def __init__(self, units: int, **kwargs) -> None:
+    def __init__(
+        self, 
+        units: int = None, 
+        normalize: bool = False,
+        skip_connection: bool = False, 
+        **kwargs
+    ) -> None:
         super().__init__(**kwargs)
         self.units = units
-        
+        self._normalize_aggregate = normalize
+        self._skip_connection = skip_connection
+    
+    def build(self, spec: tensors.GraphTensor.Spec) -> None:
+        if not self.units:
+            raise ValueError(
+                f'`self.units` needs to be a positive integer. ound: {self.units}.'
+            )
+        node_feature_dim = spec.node['feature'].shape[-1]
+        self._project_input_node_feature = (
+            self._skip_connection and (node_feature_dim != self.units)
+        )
+        if self._project_input_node_feature:
+            warn(
+                '`skip_connection` 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. '
+            )
+            self._residual_projection = self.get_dense(
+                self.units, name='residual_projection'
+            )
+        if self._normalize_aggregate:
+            self._aggregation_norm = keras.layers.LayerNormalization(
+                name='aggregation_normalization'
+            )
+            self._aggregation_norm.build([None, self.units])
+
+        super().build(spec)
+
     @abc.abstractmethod
     def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         """Compute messages.
@@ -256,206 +305,1123 @@ class GraphConv(GraphLayer):
             tensor:
                 A `GraphTensor` instance.
         """
+        
+        if self._skip_connection:
+            input_node_feature = tensor.node['feature']
+            if self._project_input_node_feature:
+                input_node_feature = self._residual_projection(input_node_feature)
+
         tensor = self.message(tensor)
         tensor = self.aggregate(tensor)
-        tensor = self.update(tensor)
-        return tensor 
 
-    def get_config(self) -> dict:
-        config = super().get_config()
-        config.update({
-            'units': self.units
-        })
-        return config
-    
-    
-@keras.saving.register_keras_serializable(package='molcraft')
-class Projection(GraphLayer):
-    """Base graph projection layer.
-    """
-    def __init__(
-        self, 
-        units: int = None, 
-        activation: str = None, 
-        field: str = 'node',
-        **kwargs
-    ) -> None:
-        super().__init__(**kwargs)
-        self.units = units
-        self._activation = keras.activations.get(activation)
-        self.field = field 
+        if self._normalize_aggregate:
+            normalized_node_feature = self._aggregation_norm(tensor.node['feature'])
+            tensor = tensor.update({'node': {'feature': normalized_node_feature}})
 
-    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
-        """Builds the layer.
-        """
-        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)
-        self._dense.build([None, feature_dim])
+        tensor = self.update(tensor)
 
-    def propagate(self, tensor: tensors.GraphTensor):
-        """Calls the layer.
-        """
-        feature = getattr(tensor, self.field)['feature']
-        feature = self._dense(feature)
-        feature = self._activation(feature)
+        if not self._skip_connection:
+            return tensor
+        
+        updated_node_feature = tensor.node['feature']
         return tensor.update(
             {
-                self.field: {
-                    'feature': feature
+                'node': {
+                    'feature': updated_node_feature + input_node_feature
                 }
             }
-        ) 
+        )
 
     def get_config(self) -> dict:
         config = super().get_config()
         config.update({
             'units': self.units,
-            'activation': keras.activations.serialize(self._activation),
-            'field': self.field,
+            'normalize': self._normalize_aggregate,
+            'skip_connection': self._skip_connection,
         })
         return config
     
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class GraphNetwork(GraphLayer):
-
-    """Graph neural network.
+class GIConv(GraphConv):
 
-    Sequentially calls graph layers (`GraphLayer`) and concatenates its output. 
-
-    Args:
-        layers (list):
-            A list of graph layers.
+    """Graph isomorphism network layer.
     """
 
-    def __init__(self, layers: list[GraphLayer], **kwargs) -> None:
-        super().__init__(**kwargs)
-        self.layers = layers
-        self._update_edge_feature = False
+    def __init__(
+        self,
+        units: int,
+        activation: keras.layers.Activation | str | None = 'relu',
+        use_bias: bool = True,
+        normalize: bool = True,
+        dropout: float = 0.0,
+        update_edge_feature: bool = True,
+        **kwargs,
+    ):
+        super().__init__(
+            units=units, 
+            normalize=normalize, 
+            use_bias=use_bias, 
+            **kwargs
+        )
+        self._activation = keras.activations.get(activation)
+        self._dropout = dropout 
+        self._update_edge_feature = update_edge_feature
 
     def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
         """Builds the layer.
         """
-        units = self.layers[0].units 
         node_feature_dim = spec.node['feature'].shape[-1]
-        if node_feature_dim != units:
-            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._update_node_feature = True 
-        has_edge_feature = 'feature' in spec.edge 
-        if has_edge_feature:
+
+        self.epsilon = self.add_weight(
+            name='epsilon', 
+            shape=(), 
+            initializer='zeros',
+            trainable=True,
+        )
+
+        self._has_edge_feature = 'feature' in spec.edge
+        if self._has_edge_feature:
             edge_feature_dim = spec.edge['feature'].shape[-1]
-            if edge_feature_dim != units:
-                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)
-                self._update_edge_feature = True
 
-    def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Calls the layer.
+            if not self._update_edge_feature:
+                if (edge_feature_dim != node_feature_dim):
+                    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.'
+                    )
+                    self._update_edge_feature = True 
+
+            if self._update_edge_feature:
+                self._edge_dense = self.get_dense(node_feature_dim)
+                self._edge_dense.build([None, edge_feature_dim])
+        else:
+            self._update_edge_feature = False
+                
+        self._feedforward_intermediate_dense = self.get_dense(self.units)
+        self._feedforward_intermediate_dense.build([None, node_feature_dim])
+
+        has_overridden_update = self.__class__.update != GIConv.update 
+        if not has_overridden_update:
+            self._feedforward_activation = self._activation
+            self._feedforward_dropout = keras.layers.Dropout(self._dropout)
+            self._feedforward_output_dense = self.get_dense(self.units)
+            self._feedforward_output_dense.build([None, self.units])
+
+        self.built = True
+
+    def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Computes messages.
         """
-        x = tensors.to_dict(tensor)
-        if self._update_node_feature:
-            x['node']['feature'] = self._node_dense(tensor.node['feature'])
+        message = tensor.gather('feature', 'source')
+        edge_feature = tensor.edge.get('feature')
         if 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'])
+            edge_feature = self._edge_dense(edge_feature)
+        if self._has_edge_feature:
+            message += edge_feature
         return tensor.update(
             {
-                'node': {
-                    'feature': keras.ops.concatenate(outputs, axis=-1)
-                }    
+                'edge': {
+                    'message': message,
+                    'feature': edge_feature
+                }
             }
         )
-    
-    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.
 
-        Args:
-            tensor (tensors.GraphTensor):
-                The graph input.
+    def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Aggregates messages.
         """
-        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 = tensor.aggregate('message')
+        node_feature += (1 + self.epsilon) * tensor.node['feature']
+        node_feature = self._feedforward_intermediate_dense(node_feature)
+        return tensor.update(
             {
                 'node': {
-                    'feature': keras.ops.concatenate(outputs, axis=-1)
+                    'feature': node_feature,
+                },
+                'edge': {
+                    'message': None,
                 }
             }
         )
-        return tensor, outputs
     
-    def get_config(self) -> dict:
-        config = super().get_config()
-        config.update(
+    def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Updates nodes. 
+        """
+        node_feature = tensor.node['feature']
+        node_feature = self._feedforward_activation(node_feature)
+        node_feature = self._feedforward_dropout(node_feature)
+        node_feature = self._feedforward_output_dense(node_feature)
+        return tensor.update(
             {
-                'layers': [
-                    keras.layers.serialize(layer) for layer in self.layers
-                ]
+                'node': {
+                    'feature': node_feature,
+                }
             }
         )
-        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)
     
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({
+            'activation': keras.activations.serialize(self._activation),
+            'dropout': self._dropout, 
+            'update_edge_feature': self._update_edge_feature
+        })
+        return config
 
-@keras.saving.register_keras_serializable(package='molcraft')
-class NodeEmbedding(GraphLayer):
 
-    """Node embedding layer.
+@keras.saving.register_keras_serializable(package='molgraphx')
+class GAConv(GraphConv):
 
-    Embeds nodes based on its initial features.
+    """Graph attention network layer.
     """
 
     def __init__(
-        self, 
-        dim: int = None, 
-        embed_context: bool = True,
-        allow_masking: bool = True, 
+        self,
+        units: int,
+        heads: int = 8,
+        activation: keras.layers.Activation | str | None = "relu",
+        use_bias: bool = True,
+        normalize: bool = True,
+        dropout: float = 0.0,
+        update_edge_feature: bool = True,
+        attention_activation: keras.layers.Activation | str | None = "leaky_relu",
+        **kwargs,
+    ) -> None:
+        kwargs['skip_connection'] = False
+        super().__init__(
+            units=units, 
+            normalize=normalize, 
+            use_bias=use_bias,
+            **kwargs
+        )
+        self._heads = heads
+        if self.units % self.heads != 0:
+            raise ValueError(f"units need to be divisible by heads.")
+        self._head_units = self.units // self.heads 
+        self._activation = keras.activations.get(activation)
+        self._dropout = dropout
+        self._normalize = normalize
+        self._update_edge_feature = update_edge_feature
+        self._attention_activation = keras.activations.get(attention_activation)
+
+    @property 
+    def heads(self):
+        return self._heads 
+    
+    @property 
+    def head_units(self):
+        return self._head_units 
+
+    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
+
+        node_feature_dim = spec.node['feature'].shape[-1]
+        attn_feature_dim = node_feature_dim + node_feature_dim
+
+        self._has_edge_feature = 'feature' in spec.edge 
+        if self._has_edge_feature:
+            edge_feature_dim = spec.edge['feature'].shape[-1]
+            attn_feature_dim += edge_feature_dim
+            if self._update_edge_feature:
+                self._edge_dense = self.get_einsum_dense(
+                    'ijh,jkh->ikh', (self.head_units, self.heads)
+                )
+                self._edge_dense.build([None, self.head_units, self.heads])
+        else:
+            self._update_edge_feature = False
+            
+        self._node_dense = self.get_einsum_dense(
+            'ij,jkh->ikh', (self.head_units, self.heads)
+        )
+        self._node_dense.build([None, node_feature_dim])
+
+        self._feature_dense = self.get_einsum_dense(
+            'ij,jkh->ikh', (self.head_units, self.heads)
+        )
+        self._feature_dense.build([None, attn_feature_dim])
+
+        self._attention_dense = self.get_einsum_dense(
+            'ijh,jkh->ikh', (1, self.heads)
+        )
+        self._attention_dense.build([None, self.head_units, self.heads])
+        
+        self._node_self_dense = self.get_einsum_dense(
+            'ij,jkh->ikh', (self.head_units, self.heads)
+        )
+        self._node_self_dense.build([None, node_feature_dim])
+        self._dropout_layer = keras.layers.Dropout(self._dropout)
+
+        self.built = True
+
+    def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+
+        attention_feature = keras.ops.concatenate(
+            [
+                tensor.gather('feature', 'source'),
+                tensor.gather('feature', 'target')
+            ], 
+            axis=-1
+        )
+        if self._has_edge_feature:
+            attention_feature = keras.ops.concatenate(
+                [
+                    attention_feature, 
+                    tensor.edge['feature']
+                ], 
+                axis=-1
+            )
+
+        attention_feature = self._feature_dense(attention_feature)
+
+        edge_feature = tensor.edge.get('feature')
+
+        if self._update_edge_feature:
+            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_score = self._attention_dense(attention_feature)
+        attention_score = ops.edge_softmax(
+            score=attention_score, edge_target=tensor.edge['target']
+        )
+        node_feature = self._node_dense(tensor.node['feature'])
+        message = ops.gather(node_feature, tensor.edge['source'])
+        return tensor.update(
+            {
+                'edge': {
+                    'message': message,
+                    'weight': attention_score,
+                    'feature': edge_feature,
+                }
+            }
+        )
+    
+    def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        node_feature = tensor.aggregate('message')
+        node_feature += self._node_self_dense(tensor.node['feature'])
+        node_feature = self._dropout_layer(node_feature)
+        node_feature = keras.ops.reshape(node_feature, (-1, self.units))
+        return tensor.update(
+            {
+                'node': {
+                    'feature': node_feature
+                },
+                'edge': {
+                    'message': None,
+                    'weight': None,
+                }
+            }
+        )
+    
+    def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        node_feature = self._activation(tensor.node['feature'])
+        return tensor.update(
+            {
+                'node': {
+                    'feature': node_feature
+                }
+            }
+        )
+    
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({
+            "heads": self._heads,
+            'activation': keras.activations.serialize(self._activation),
+            'dropout': self._dropout, 
+            'update_edge_feature': self._update_edge_feature,
+            'attention_activation': keras.activations.serialize(self._attention_activation),
+        })
+        return config
+    
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class GTConv(GraphConv):
+
+    """Graph transformer layer.
+    """
+
+    def __init__(
+        self,
+        units: int,
+        heads: int = 8,
+        activation: keras.layers.Activation | str | None = "relu",
+        use_bias: bool = True,
+        normalize: bool = True,
+        dropout: float = 0.0,
+        attention_dropout: float = 0.0,
+        **kwargs,
+    ) -> None:
+        kwargs['skip_connection'] = False
+        super().__init__(
+            units=units, 
+            normalize=normalize, 
+            use_bias=use_bias,
+            **kwargs
+        )
+        self._heads = heads
+        if self.units % self.heads != 0:
+            raise ValueError(f"units need to be divisible by heads.")
+        self._head_units = self.units // self.heads 
+        self._activation = keras.activations.get(activation)
+        self._dropout = dropout
+        self._attention_dropout = attention_dropout
+        self._normalize = normalize
+
+    @property 
+    def heads(self):
+        return self._heads 
+    
+    @property 
+    def head_units(self):
+        return self._head_units 
+    
+    def build_from_spec(self, spec):
+        """Builds the layer.
+        """
+        node_feature_dim = spec.node['feature'].shape[-1]
+        self.project_residual = node_feature_dim != self.units
+        if self.project_residual:
+            warn(
+                '`GTConv` uses residual connections, 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. '
+            )   
+            self._residual_dense = self.get_dense(self.units)
+            self._residual_dense.build([None, node_feature_dim])
+            
+        self._query_dense = self.get_einsum_dense(
+            'ij,jkh->ikh', (self.head_units, self.heads)
+        )
+        self._query_dense.build([None, node_feature_dim])
+
+        self._key_dense = self.get_einsum_dense(
+            'ij,jkh->ikh', (self.head_units, self.heads)
+        )
+        self._key_dense.build([None, node_feature_dim])
+
+        self._value_dense = self.get_einsum_dense(
+            'ij,jkh->ikh', (self.head_units, self.heads)
+        )
+        self._value_dense.build([None, node_feature_dim])
+
+        self._output_dense = self.get_dense(self.units)
+        self._output_dense.build([None, self.units])
+
+        self._softmax_dropout = keras.layers.Dropout(self._attention_dropout) 
+
+        self._self_attention_dropout = keras.layers.Dropout(self._dropout)
+
+        self._add_bias = not 'bias' in spec.edge
+
+        if self._add_bias:
+            self._edge_bias = EdgeBias(biases=self.heads)
+            self._edge_bias.build_from_spec(spec)
+
+        has_overridden_update = self.__class__.update != GTConv.update 
+        if not has_overridden_update:
+            
+            if self._normalize:
+                self._feedforward_output_norm = keras.layers.LayerNormalization()
+                self._feedforward_output_norm.build([None, self.units])
+
+            self._feedforward_dropout = keras.layers.Dropout(self._dropout)
+
+            self._feedforward_intermediate_dense = self.get_dense(self.units)
+            self._feedforward_intermediate_dense.build([None, self.units])
+
+            self._feedforward_output_dense = self.get_dense(self.units)
+            self._feedforward_output_dense.build([None, self.units])
+
+        self.built = True
+
+    def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Computes messages.
+        """
+        if self._add_bias:
+            edge_bias = self._edge_bias(tensor)
+            tensor = tensor.update(
+                {
+                    'edge': {
+                        'bias': edge_bias
+                    }
+                }
+            )
+            
+        node_feature = tensor.node['feature']
+        
+        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)
+
+        return tensor.update(
+            {
+                'edge': {
+                    'message': value,
+                    'weight': attention,
+                },
+            }
+        )
+
+    def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Aggregates messages.
+        """
+        node_feature = tensor.aggregate('message')
+        node_feature = keras.ops.reshape(node_feature, (-1, self.units))
+        node_feature = self._output_dense(node_feature)
+        node_feature = self._self_attention_dropout(node_feature)
+        return tensor.update(
+            {
+                'node': {
+                    'feature': node_feature,
+                    'residual': tensor.node['feature']
+                },
+                'edge': {
+                    'message': None,
+                    'weight': None,
+                }
+            }
+        )
+
+    def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Updates nodes. 
+        """
+        node_feature = tensor.node['feature']
+
+        residual = tensor.node['residual']
+        if self.project_residual:
+            residual = self._residual_dense(residual)
+
+        node_feature += residual
+        residual = node_feature
+
+        node_feature = self._feedforward_intermediate_dense(node_feature)
+        node_feature = self._activation(node_feature)
+        node_feature = self._feedforward_output_dense(node_feature)
+        node_feature = self._feedforward_dropout(node_feature)
+        if self._normalize:
+            node_feature = self._feedforward_output_norm(node_feature)
+
+        node_feature += residual
+
+        return tensor.update(
+            {
+                'node': {
+                    'feature': node_feature,
+                },
+            }
+        )
+    
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({
+            "heads": self._heads,
+            'activation': keras.activations.serialize(self._activation),
+            'dropout': self._dropout, 
+            'attention_dropout': self._attention_dropout,
+        })
+        return config
+    
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class GTConv3D(GTConv):
+
+    """Graph transformer 3D layer.
+    """
+
+    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
+        super().build_from_spec(spec)
+        if self._add_bias:
+            node_feature_dim = spec.node['feature'].shape[-1]
+            kernels = self.units
+            self._gaussian_basis = GaussianDistance(kernels)
+            self._gaussian_basis.build_from_spec(spec)
+            self._centrality_dense = self.get_dense(units=node_feature_dim)
+            self._centrality_dense.build([None, kernels])
+            self._gaussian_edge_bias = self.get_dense(self.heads)
+            self._gaussian_edge_bias.build([None, kernels])
+
+    def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Computes messages.
+        """
+        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)
+        
+        return tensor.update(
+            {
+                'edge': {
+                    'message': value,
+                    'weight': attention,
+                },
+            }
+        )
+
+    def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Aggregates messages.
+        """
+        node_feature = tensor.aggregate('message')
+        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)
+        node_feature = self._self_attention_dropout(node_feature)
+        return tensor.update(
+            {
+                'node': {
+                    'feature': node_feature,
+                    'residual': tensor.node['feature']
+                },
+                'edge': {
+                    'message': None,
+                    'weight': None,
+                }
+            }
+        )
+    
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class MPConv(GraphConv):
+
+    """Message passing neural network layer.
+    """
+
+    def __init__(
+        self, 
+        units: int = 128, 
+        activation: keras.layers.Activation | str | None = None, 
+        use_bias: bool = True,
+        normalize: bool = True,
+        dropout: float = 0.0,
+        **kwargs
+    ) -> None:
+        super().__init__(
+            units=units, 
+            normalize=normalize, 
+            use_bias=use_bias,
+            **kwargs
+        )
+        self._activation = keras.activations.get(activation)
+        self._dropout = dropout or 0.0
+
+    def build_from_spec(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)
+        self.built = True
+
+    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')
+        previous = tensor.node['feature']
+        if self.project_input_node_feature:
+            previous = self._previous_node_dense(previous)
+        return tensor.update(
+            {
+                'node': {
+                    'feature': aggregate,
+                    'previous_feature': previous,
+                }
+            }
+        )
+
+    def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        updated_node_feature, _ = self.update_fn(
+            inputs=tensor.node['feature'],
+            states=tensor.node['previous_feature']
+        )
+        return tensor.update(
+            {
+                'node': {
+                    'feature': updated_node_feature,
+                    'previous_feature': None,
+                }
+            }
+        )
+
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({
+            'activation': keras.activations.serialize(self._activation),
+            'dropout': self._dropout,
+        })
+        return config 
+    
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class MPConv3D(MPConv):
+
+    """3D Message passing neural network layer.
+    """
+    
+    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.
+    """
+
+    def __init__(
+        self, 
+        units: int = 128, 
+        activation: keras.layers.Activation | str | None = None, 
+        use_bias: bool = True,
+        normalize: bool = True,
+        dropout: float = 0.0,
+        **kwargs
+    ) -> None:
+        super().__init__(
+            units=units, 
+            normalize=normalize, 
+            use_bias=use_bias,
+            **kwargs
+        )
+        self._activation = keras.activations.get(activation)
+        self._dropout = dropout or 0.0
+
+    def build_from_spec(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.'
+            )
+        node_feature_dim = spec.node['feature'].shape[-1]
+        feature_dim = node_feature_dim + node_feature_dim + 1
+        if 'feature' in spec.edge:
+            self._has_edge_feature = True
+            edge_feature_dim = spec.edge['feature'].shape[-1]
+            feature_dim += edge_feature_dim
+        else:
+            self._has_edge_feature = False
+
+        self.message_fn = self.get_dense(self.units, activation=self._activation)
+        self.message_fn.build([None, feature_dim])
+        self.dense_position = self.get_dense(1)
+        self.dense_position.build([None, self.units])
+
+        has_overridden_update = self.__class__.update != EGConv3D.update 
+        if not has_overridden_update:
+            self.update_fn = self.get_dense(self.units, activation=self._activation)
+            self.update_fn.build([None, node_feature_dim + self.units])
+            self._dropout_layer = keras.layers.Dropout(self._dropout)
+        self.built = True
+
+    def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Computes messages.
+        """
+        relative_node_coordinate = keras.ops.subtract(
+            tensor.gather('coordinate', 'target'), 
+            tensor.gather('coordinate', 'source')
+        ) 
+        euclidean_distance = keras.ops.sum(
+            keras.ops.square(
+                relative_node_coordinate
+            ), 
+            axis=-1, 
+            keepdims=True
+        )
+        feature = keras.ops.concatenate(
+            [
+                tensor.gather('feature', 'target'), 
+                tensor.gather('feature', 'source'), 
+                euclidean_distance, 
+            ], 
+            axis=-1
+        )
+        if self._has_edge_feature:
+            feature = keras.ops.concatenate(
+                [
+                    feature,
+                    tensor.edge['feature']
+                ], 
+                axis=-1
+            )
+        message = self.message_fn(feature)
+        relative_node_coordinate = keras.ops.multiply(
+            relative_node_coordinate,
+            self.dense_position(message)
+        )
+        return tensor.update(
+            {
+                'edge': {
+                    'message': message,
+                    'relative_node_coordinate': relative_node_coordinate
+                }
+            }
+        )
+
+    def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Aggregates messages.
+        """
+        coefficient = keras.ops.bincount(
+            tensor.edge['source'], 
+            minlength=tensor.num_nodes
+        )
+        coefficient = keras.ops.cast(
+            coefficient, tensor.node['coordinate'].dtype
+        )
+        coefficient = keras.ops.expand_dims(
+            keras.ops.divide_no_nan(1, coefficient), axis=1
+        )  
+
+        updated_coordinate = tensor.aggregate('relative_node_coordinate') * coefficient
+        updated_coordinate += tensor.node['coordinate']
+
+        aggregate = tensor.aggregate('message')
+        return tensor.update(
+            {
+                'node': {
+                    'feature': aggregate, 
+                    'coordinate': updated_coordinate,
+                    'previous_feature': tensor.node['feature'],
+                },
+                'edge': {
+                    'message': None,
+                    'relative_node_coordinate': None
+                }
+            }
+        ) 
+
+    def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Updates nodes.
+        """
+        updated_node_feature = self.update_fn(
+            keras.ops.concatenate(
+                [
+                    tensor.node['feature'], 
+                    tensor.node['previous_feature']
+                ], 
+                axis=-1
+            )
+        )  
+        updated_node_feature = self._dropout_layer(updated_node_feature)
+        return tensor.update(
+            {
+                'node': {
+                    'feature': updated_node_feature, 
+                    'previous_feature': None,
+                },
+            }
+        )
+    
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({
+            'activation': keras.activations.serialize(self._activation),
+            'dropout': self._dropout,
+        })
+        return config 
+    
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class Projection(GraphLayer):
+    """Base graph projection layer.
+    """
+    def __init__(
+        self, 
+        units: int = None, 
+        activation: str = None, 
+        field: str = 'node',
+        **kwargs
+    ) -> None:
+        super().__init__(**kwargs)
+        self.units = units
+        self._activation = keras.activations.get(activation)
+        self.field = field 
+
+    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
+        """Builds the layer.
+        """
+        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)
+        self._dense.build([None, feature_dim])
+        self.built = True
+
+    def propagate(self, tensor: tensors.GraphTensor):
+        """Calls the layer.
+        """
+        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 GraphNetwork(GraphLayer):
+
+    """Graph neural network.
+
+    Sequentially calls graph layers (`GraphLayer`) and concatenates its output. 
+
+    Args:
+        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_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
+        """Builds the layer.
+        """
+        units = self.layers[0].units 
+        node_feature_dim = spec.node['feature'].shape[-1]
+        if node_feature_dim != units:
+            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._update_node_feature = True 
+        has_edge_feature = 'feature' in spec.edge 
+        if has_edge_feature:
+            edge_feature_dim = spec.edge['feature'].shape[-1]
+            if edge_feature_dim != units:
+                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)
+                self._update_edge_feature = True
+        self.built = True
+
+    def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Calls the layer.
+        """
+        x = tensors.to_dict(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'])
+        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.
+
+        Args:
+            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 NodeEmbedding(GraphLayer):
+
+    """Node embedding layer.
+
+    Embeds nodes based on its initial features.
+    """
+
+    def __init__(
+        self, 
+        dim: int = None, 
+        normalize: bool = True,
+        embed_context: bool = True,
+        allow_masking: bool = True, 
         **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
@@ -482,6 +1448,12 @@ class NodeEmbedding(GraphLayer):
             context_feature_dim = spec.context['feature'].shape[-1]
             self._context_dense = self.get_dense(self.dim)
             self._context_dense.build([None, context_feature_dim])
+        
+        if self._normalize:
+            self._norm = keras.layers.LayerNormalization()
+            self._norm.build([None, self.dim])
+
+        self.built = True
 
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         """Calls the layer.
@@ -515,6 +1487,9 @@ class NodeEmbedding(GraphLayer):
             # Slience warning of 'no gradients for variables'
             feature = feature + (self._mask_feature * 0.0)
 
+        if self._normalize:
+            feature = self._norm(feature)
+
         return tensor.update({'node': {'feature': feature}})
 
     @property 
@@ -534,6 +1509,8 @@ class NodeEmbedding(GraphLayer):
         config = super().get_config()
         config.update({
             'dim': self.dim,
+            'normalize': self._normalize,
+            'embed_context': self._embed_context,
             'allow_masking': self._allow_masking
         })
         return config
@@ -544,503 +1521,210 @@ class EdgeEmbedding(GraphLayer):
 
     """Edge embedding layer.
 
-    Embeds edges based on its initial features.
-    """
-
-    def __init__(
-        self, 
-        dim: int = None, 
-        allow_masking: bool = True, 
-        **kwargs
-    ) -> None:
-        super().__init__(**kwargs)
-        self.dim = dim
-        self._masking_rate = None
-        self._allow_masking = allow_masking
-
-    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
-        """Builds the layer.
-        """
-        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.build([None, feature_dim])
-
-        self._has_super = 'super' in spec.edge
-        if self._has_super:
-            self._super_feature = self.get_weight(shape=[self.dim], name='super_edge_feature')
-        if self._allow_masking:
-            self._mask_feature = self.get_weight(shape=[self.dim], name='mask_edge_feature')
-
-    def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Calls the layer.
-        """
-        feature = self._edge_dense(tensor.edge['feature'])
-
-        if self._has_super:
-            super_feature = self._super_feature
-            super_mask = keras.ops.expand_dims(tensor.edge['super'], 1)
-            feature = keras.ops.where(super_mask, super_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:
-            # Slience warning of 'no gradients for variables'
-            feature = feature + (self._mask_feature * 0.0)
-
-        return tensor.update({'edge': {'feature': 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)
-
-    def get_config(self) -> dict:
-        config = super().get_config()
-        config.update({
-            'dim': self.dim,
-            'allow_masking': self._allow_masking
-        })
-        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):
-        super().__init__(units=units, activation=activation, field='context', **kwargs)
-
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class NodeProjection(Projection):
-    """Node projection layer.
-    """
-    def __init__(self, units: int = None, activation: str = None, **kwargs):
-        super().__init__(units=units, activation=activation, field='node', **kwargs)
-
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class EdgeProjection(Projection):
-    """Edge projection layer.
-    """
-    def __init__(self, units: int = None, activation: str = None, **kwargs):
-        super().__init__(units=units, activation=activation, field='edge', **kwargs)
-
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class GINConv(GraphConv):
-
-    """Graph isomorphism network layer.
-    """
-
-    def __init__(
-        self,
-        units: int,
-        activation: keras.layers.Activation | str | None = 'relu',
-        dropout: float = 0.0,
-        normalize: bool = True,
-        update_edge_feature: bool = True,
-        **kwargs,
-    ):
-        super().__init__(units=units, **kwargs)
-        self._activation = keras.activations.get(activation)
-        self._normalize = normalize
-        self._dropout = dropout 
-        self._update_edge_feature = update_edge_feature
-
-    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
-        """Builds the layer.
-        """
-        node_feature_dim = spec.node['feature'].shape[-1]
-
-        self.epsilon = self.add_weight(
-            name='epsilon', 
-            shape=(), 
-            initializer='zeros',
-            trainable=True,
-        )
-
-        if 'feature' in spec.edge:
-            edge_feature_dim = spec.edge['feature'].shape[-1]
-
-            if not self._update_edge_feature:
-                if (edge_feature_dim != node_feature_dim):
-                    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.'
-                    )
-                    self._update_edge_feature = True 
-
-            if self._update_edge_feature:
-                self._edge_dense = self.get_dense(node_feature_dim)
-                self._edge_dense.build([None, edge_feature_dim])
-        else:
-            self._update_edge_feature = False
-                
-        has_overridden_update = self.__class__.update != GINConv.update 
-        if not has_overridden_update:
-            # Use default feedforward network
-            self._feedforward_intermediate_dense = self.get_dense(self.units)
-            self._feedforward_intermediate_dense.build([None, node_feature_dim])
-
-            if self._normalize:
-                self._feedforward_intermediate_norm = keras.layers.BatchNormalization()
-                self._feedforward_intermediate_norm.build([None, self.units])
-
-            self._feedforward_dropout = keras.layers.Dropout(self._dropout)
-            self._feedforward_activation = self._activation
-                
-            self._feedforward_output_dense = self.get_dense(self.units)
-            self._feedforward_output_dense.build([None, self.units])
-    
-    def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Compute messages.
-        """
-        message = tensor.gather('feature', 'source')
-        edge_feature = tensor.edge.get('feature')
-        if self._update_edge_feature:
-            edge_feature = self._edge_dense(edge_feature)
-        if edge_feature is not None:
-            message += edge_feature
-        return tensor.update(
-            {
-                'edge': {
-                    'message': message,
-                    'feature': edge_feature
-                }
-            }
-        )
-
-    def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Aggregates messages.
-        """
-        node_feature = tensor.aggregate('message')
-        node_feature += (1 + self.epsilon) * tensor.node['feature']
-        return tensor.update(
-            {
-                'node': {
-                    'feature': node_feature,
-                },
-                'edge': {
-                    'message': None,
-                }
-            }
-        )
-    
-    def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Updates nodes. 
-        """
-        node_feature = tensor.node['feature']
-        node_feature = self._feedforward_intermediate_dense(node_feature)
-        node_feature = self._feedforward_activation(node_feature)
-        if self._normalize:
-            node_feature = self._feedforward_intermediate_norm(node_feature)
-        node_feature = self._feedforward_dropout(node_feature)
-        node_feature = self._feedforward_output_dense(node_feature)
-        return tensor.update(
-            {
-                'node': {
-                    'feature': node_feature,
-                }
-            }
-        )
-    
-    def get_config(self) -> dict:
-        config = super().get_config()
-        config.update({
-            'activation': keras.activations.serialize(self._activation),
-            'dropout': self._dropout, 
-            'normalize': self._normalize,
-        })
-        return config
-
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class GTConv(GraphConv):
-
-    """Graph transformer layer.
+    Embeds edges based on its initial features.
     """
 
     def __init__(
-        self,
-        units: int,
-        heads: int = 8,
-        activation: keras.layers.Activation | str | None = "relu",
-        dropout: float = 0.0,
-        attention_dropout: float = 0.0,
+        self, 
+        dim: int = None, 
         normalize: bool = True,
-        normalize_first: bool = True,
-        **kwargs,
+        allow_masking: bool = True, 
+        **kwargs
     ) -> None:
-        super().__init__(units=units, **kwargs)
-        self._heads = heads
-        if self.units % self.heads != 0:
-            raise ValueError(f"units need to be divisible by heads.")
-        self._head_units = self.units // self.heads 
-        self._activation = keras.activations.get(activation)
-        self._dropout = dropout
-        self._attention_dropout = attention_dropout
+        super().__init__(**kwargs)
+        self.dim = dim
         self._normalize = normalize
-        self._normalize_first = normalize_first
+        self._masking_rate = None
+        self._allow_masking = allow_masking
 
-    @property 
-    def heads(self):
-        return self._heads 
-    
-    @property 
-    def head_units(self):
-        return self._head_units 
-    
-    def build_from_spec(self, spec):
+    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
         """Builds the layer.
         """
-        node_feature_dim = spec.node['feature'].shape[-1]
-        incompatible_dim = node_feature_dim != self.units
-        if incompatible_dim:
-            warnings.warn(
-                message=(
-                    '`GTConv` uses residual connections, but input node feature dim '
-                    'is incompatible with intermediate dim (`units`). '
-                    'Automatically projecting first residual to match its dim with intermediate dim.'
-                ),
-                category=UserWarning,
-                stacklevel=1
-            )
-            self._residual_dense = self.get_dense(self.units)
-            self._residual_dense.build([None, node_feature_dim])
-            self._project_residual = True
-        else:
-            self._project_residual = False
-            
-        self._query_dense = self.get_einsum_dense(
-            'ij,jkh->ikh', (self.head_units, self.heads)
-        )
-        self._query_dense.build([None, node_feature_dim])
-
-        self._key_dense = self.get_einsum_dense(
-            'ij,jkh->ikh', (self.head_units, self.heads)
-        )
-        self._key_dense.build([None, node_feature_dim])
-
-        self._value_dense = self.get_einsum_dense(
-            'ij,jkh->ikh', (self.head_units, self.heads)
-        )
-        self._value_dense.build([None, node_feature_dim])
-
-        self._output_dense = self.get_dense(self.units)
-        self._output_dense.build([None, self.units])
+        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.build([None, feature_dim])
 
-        self._softmax_dropout = keras.layers.Dropout(self._attention_dropout) 
+        self._has_super = 'super' in spec.edge
+        if self._has_super:
+            self._super_feature = self.get_weight(shape=[self.dim], name='super_edge_feature')
+        if self._allow_masking:
+            self._mask_feature = self.get_weight(shape=[self.dim], name='mask_edge_feature')
+        if self._normalize:
+            self._norm = keras.layers.LayerNormalization()
+            self._norm.build([None, self.dim])
 
-        self._self_attention_norm = keras.layers.LayerNormalization()
-        if self._normalize_first:
-            self._self_attention_norm.build([None, node_feature_dim])
-        else:
-            self._self_attention_norm.build([None, self.units])
+        self.built = True
 
-        self._self_attention_dropout = keras.layers.Dropout(self._dropout)
+    def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Calls the layer.
+        """
+        feature = self._edge_dense(tensor.edge['feature'])
 
-        has_overriden_edge_bias = (
-            self.__class__.add_edge_bias != GTConv.add_edge_bias
-        )
-        if not has_overriden_edge_bias:
-            self._has_edge_length = 'length' in spec.edge
-            if self._has_edge_length and 'bias' not in spec.edge:
-                edge_length_dim = spec.edge['length'].shape[-1]
-                self._spatial_encoding_dense = self.get_einsum_dense(
-                    'ij,jkh->ikh', (1, self.heads), kernel_initializer='zeros'
-                )
-                self._spatial_encoding_dense.build([None, edge_length_dim])
+        if self._has_super:
+            super_feature = self._super_feature
+            super_mask = keras.ops.expand_dims(tensor.edge['super'], 1)
+            feature = keras.ops.where(super_mask, super_feature, feature)
 
-            self._has_edge_feature = 'feature' in spec.edge
-            if self._has_edge_feature and 'bias' not in spec.edge:
-                edge_feature_dim = spec.edge['feature'].shape[-1]
-                self._edge_feature_dense = self.get_einsum_dense(
-                    'ij,jkh->ikh', (1, self.heads),
+        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'])
                 )
-                self._edge_feature_dense.build([None, edge_feature_dim])
-
-        has_overridden_update = self.__class__.update != GTConv.update 
-        if not has_overridden_update:
-
-            self._feedforward_norm = keras.layers.LayerNormalization()
-            self._feedforward_norm.build([None, self.units])
-
-            self._feedforward_dropout = keras.layers.Dropout(self._dropout)
+            mask = keras.ops.expand_dims(mask, -1)
+            feature = keras.ops.where(mask, self._mask_feature, feature)
+        elif self._allow_masking:
+            # Slience warning of 'no gradients for variables'
+            feature = feature + (self._mask_feature * 0.0)
 
-            self._feedforward_intermediate_dense = self.get_dense(self.units)
-            self._feedforward_intermediate_dense.build([None, self.units])
+        if self._normalize:
+            feature = self._norm(feature)
 
-            self._feedforward_output_dense = self.get_dense(self.units)
-            self._feedforward_output_dense.build([None, self.units])
+        return tensor.update({'edge': {'feature': feature}})
 
-    def add_node_bias(self, tensor: tensors.GraphTensor) -> tf.Tensor:
-        return tensor 
+    @property 
+    def masking_rate(self):
+        return self._masking_rate 
     
-    def add_edge_bias(self, tensor: tensors.GraphTensor) -> tf.Tensor:
-        if 'bias' in tensor.edge:
-            return tensor 
-        elif not self._has_edge_feature and not self._has_edge_length:
-            return tensor 
-        
-        if self._has_edge_feature and not self._has_edge_length:
-            edge_bias = self._edge_feature_dense(tensor.edge['feature'])
-        elif not self._has_edge_feature and self._has_edge_length:
-            edge_bias = self._spatial_encoding_dense(tensor.edge['length'])
-        else:
-            edge_bias = (
-                self._edge_feature_dense(tensor.edge['feature']) + 
-                self._spatial_encoding_dense(tensor.edge['length'])
+    @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`.'
             )
- 
-        return tensor.update(
-            {
-                'edge': {
-                    'bias': edge_bias
-                }
-            }
-        )
-
-    def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Compute messages.
-        """
-        tensor = self.add_edge_bias(tensor)
-        tensor = self.add_node_bias(tensor)
+        self._masking_rate = float(rate)
 
-        node_feature = tensor.node['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
+    
 
-        if 'bias' in tensor.node:
-            node_feature += tensor.node['bias']
-        
-        if self._normalize_first:
-            node_feature = self._self_attention_norm(node_feature)
-        
-        query = self._query_dense(node_feature)
-        key = self._key_dense(node_feature)
-        value = self._value_dense(node_feature)
+@keras.saving.register_keras_serializable(package='molcraft')
+class ContextProjection(Projection):
+    """Context projection layer.
+    """
+    def __init__(self, units: int = None, activation: str = None, **kwargs):
+        super().__init__(units=units, activation=activation, field='context', **kwargs)
 
-        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.units))
+@keras.saving.register_keras_serializable(package='molcraft')
+class NodeProjection(Projection):
+    """Node projection layer.
+    """
+    def __init__(self, units: int = None, activation: str = None, **kwargs):
+        super().__init__(units=units, activation=activation, field='node', **kwargs)
 
-        if 'bias' in tensor.edge:
-            attention_score += tensor.edge['bias']
-        
-        attention = ops.edge_softmax(attention_score, tensor.edge['target'])
-        attention = self._softmax_dropout(attention)
 
-        return tensor.update(
-            {
-                'edge': {
-                    'message': value,
-                    'weight': attention,
-                },
-            }
-        )
+@keras.saving.register_keras_serializable(package='molcraft')
+class EdgeProjection(Projection):
+    """Edge projection layer.
+    """
+    def __init__(self, units: int = None, activation: str = None, **kwargs):
+        super().__init__(units=units, activation=activation, field='edge', **kwargs)
 
-    def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Aggregates messages.
-        """
-        node_feature = tensor.aggregate('message')
 
-        node_feature = keras.ops.reshape(node_feature, (-1, self.units))
-        node_feature = self._output_dense(node_feature)
-        node_feature = self._self_attention_dropout(node_feature)
+@keras.saving.register_keras_serializable(package='molcraft')
+class EdgeBias(GraphLayer):
 
-        residual = tensor.node['feature']
-        if self._project_residual:
-            residual = self._residual_dense(residual)
-        node_feature += residual
+    def __init__(self, biases: int, **kwargs):
+        super().__init__(**kwargs)
+        self.biases = biases
 
-        if not self._normalize_first:
-            node_feature = self._self_attention_norm(node_feature)
+    def build_from_spec(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)
+            self._edge_feature_dense.build([None, spec.edge['feature'].shape[-1]])
+        if self._has_edge_length:
+            self._edge_length_dense = self.get_dense(
+                self.biases, kernel_initializer='zeros'
+            )
+            self._edge_length_dense.build([None, spec.edge['length'].shape[-1]])
+        self.built = True
 
-        return tensor.update(
-            {
-                'node': {
-                    'feature': node_feature,
-                },
-                'edge': {
-                    'message': None,
-                    'weight': None,
-                }
-            }
+    def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        bias = keras.ops.zeros(
+            shape=(tensor.num_edges, self.biases), 
+            dtype=tensor.node['feature'].dtype
         )
-    
-
-    def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Updates nodes. 
-        """
-        node_feature = tensor.node['feature']
-
-        if self._normalize_first:
-            node_feature = self._feedforward_norm(node_feature)
+        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
 
-        node_feature = self._feedforward_intermediate_dense(node_feature)
-        node_feature = self._activation(node_feature)
-        node_feature = self._feedforward_output_dense(node_feature)
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({'biases': self.biases})
+        return config
+    
 
-        node_feature = self._feedforward_dropout(node_feature)
-        node_feature += tensor.node['feature']
+@keras.saving.register_keras_serializable(package='molcraft')
+class GaussianDistance(GraphLayer):
 
-        if not self._normalize_first:
-            node_feature = self._feedforward_norm(node_feature)
+    def __init__(self, kernels: int, **kwargs):
+        super().__init__(**kwargs)
+        self.kernels = kernels 
 
-        return tensor.update(
-            {
-                'node': {
-                    'feature': node_feature,
-                },
-            }
+    def build_from_spec(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
         )
-    
+        self.built = True
+        
+    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
+        )
+
     def get_config(self) -> dict:
         config = super().get_config()
         config.update({
-            "heads": self._heads,
-            'activation': keras.activations.serialize(self._activation),
-            'dropout': self._dropout, 
-            'attention_dropout': self._attention_dropout,
-            'normalize': self._normalize,
-            'normalize_first': self._normalize_first,
+            'kernels': self.kernels,
         })
         return config
-    
+
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class Readout(keras.layers.Layer):
+class Readout(GraphLayer):
+
+    """Readout layer.
+    """
 
     def __init__(self, mode: str | None = None, **kwargs):
+        kwargs['kernel_initializer'] = None 
+        kwargs['bias_initializer'] = None
         super().__init__(**kwargs)
         self.mode = mode
-        if not self.mode:
-            self._reduce_fn = None
-        elif str(self.mode).lower().startswith('sum'):
+        if str(self.mode).lower().startswith('sum'):
             self._reduce_fn = keras.ops.segment_sum
         elif str(self.mode).lower().startswith('max'):
             self._reduce_fn = keras.ops.segment_max 
@@ -1052,50 +1736,25 @@ class Readout(keras.layers.Layer):
     def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
         """Builds the layer.
         """
-        pass
+        self.built = True
 
-    def reduce(self, tensor: tensors.GraphTensor) -> tf.Tensor:
-        if self._reduce_fn is None:
-            raise NotImplementedError("Need to define a reduce method.")
+    def propagate(self, tensor: tensors.GraphTensor) -> tf.Tensor:
+        """Calls the layer.
+        """
+        node_feature = tensor.node['feature']
         if str(self.mode).lower().startswith('super'):
             node_feature = keras.ops.where(
-                tensor.node['super'][:, None], tensor.node['feature'], 0.0
-            )
-            return self._reduce_fn(
-                node_feature, tensor.graph_indicator, tensor.num_subgraphs
+                tensor.node['super'][:, None], node_feature, 0.0
             )
         return self._reduce_fn(
-            tensor.node['feature'], tensor.graph_indicator, tensor.num_subgraphs
+            node_feature, tensor.graph_indicator, tensor.num_subgraphs
         )
 
-    def build(self, input_shapes) -> None:
-        spec = tensors.GraphTensor.Spec.from_input_shape_dict(input_shapes)
-        self.build_from_spec(spec)
-        self.built = True 
-
-    def call(self, graph) -> tf.Tensor:
-        graph_tensor = tensors.from_dict(graph) 
-        if tensors.is_ragged(graph_tensor):
-            graph_tensor = graph_tensor.flatten()
-        return self.reduce(graph_tensor)
-
-    def __call__(
-        self, 
-        graph: tensors.GraphTensor, 
-        *args, 
-        **kwargs
-    ) -> tensors.GraphTensor:
-        is_tensor = isinstance(graph, tensors.GraphTensor)
-        if is_tensor:
-            graph = tensors.to_dict(graph)
-        tensor = super().__call__(graph, *args, **kwargs)
-        return tensor
-    
     def get_config(self) -> dict:
         config = super().get_config()
         config['mode'] = self.mode 
         return config 
-
+    
 
 def Input(spec: tensors.GraphTensor.Spec) -> dict:
     """Used to specify inputs to model.
@@ -1212,13 +1871,6 @@ def _spec_from_inputs(inputs):
     return tensors.GraphTensor.Spec(**nested_specs)
 
 
-GraphTransformer = GTConvolution = GTConv
-GINConvolution = GINConv
-
-EdgeEmbed = EdgeEmbedding
-NodeEmbed = NodeEmbedding
-
-ContextDense = ContextProjection
-EdgeDense = EdgeProjection
-NodeDense = NodeProjection
+GraphTransformer = GTConv
+GraphTransformer3D = GTConv3D