molcraft 0.1.0a2__py3-none-any.whl → 0.1.0a4__py3-none-any.whl

molcraft/layers.py

@@ -1,7 +1,7 @@
-import abc
 import keras 
 import tensorflow as tf
 import warnings
+import functools
 from keras.src.models import functional
 
 from molcraft import tensors
@@ -12,11 +12,39 @@ from molcraft import ops
 class GraphLayer(keras.layers.Layer):
     """Base graph layer.
 
-    Currently, the `GraphLayer` only supports `GraphTensor` input.
+    Subclasses must implement a forward pass via **propagate(graph)**. 
+    
+    Subclasses may create dense layers and weights in **build(graph_spec)**.
+
+    Note: `GraphLayer` currently only supports `GraphTensor` input.
 
-    The list of arguments are only relevant if the derived layer 
+    The list of arguments below are only relevant if the derived layer 
     invokes 'get_dense_kwargs`, `get_dense`  or `get_einsum_dense`. 
 
+    Arguments:
+        use_bias (bool):
+            Whether bias should be used in dense layers. Default to `True`.
+        kernel_initializer (keras.initializers.Initializer, str):
+            Initializer for the kernel weight matrix of the dense layers.
+            Default to `glorot_uniform`.
+        bias_initializer (keras.initializers.Initializer, str):
+            Initializer for the bias weight vector of the dense layers.
+            Default to `zeros`.
+        kernel_regularizer (keras.regularizers.Regularizer, None):
+            Regularizer function applied to the kernel weight matrix.
+            Default to `None`.
+        bias_regularizer (keras.regularizers.Regularizer, None):
+            Regularizer function applied to the bias weight vector.
+            Default to `None`.
+        activity_regularizer (keras.regularizers.Regularizer, None):
+            Regularizer function applied to the output of the dense layers.
+            Default to `None`.
+        kernel_constraint (keras.constraints.Constraint, None):
+            Constraint function applied to the kernel weight matrix.
+            Default to `None`.
+        bias_constraint (keras.constraints.Constraint, None):
+            Constraint function applied to the bias weight vector.
+            Default to `None`.
     """
 
     def __init__(
@@ -31,73 +59,61 @@ class GraphLayer(keras.layers.Layer):
         bias_constraint: keras.constraints.Constraint | None = None,
         **kwargs,
     ) -> None:
-        super().__init__(activity_regularizer=activity_regularizer, **kwargs)
+        super().__init__(**kwargs)
         self._use_bias = use_bias
         self._kernel_initializer = keras.initializers.get(kernel_initializer)
         self._bias_initializer = keras.initializers.get(bias_initializer)
         self._kernel_regularizer = keras.regularizers.get(kernel_regularizer)
         self._bias_regularizer = keras.regularizers.get(bias_regularizer)
+        self._activity_regularizer = keras.regularizers.get(activity_regularizer)
         self._kernel_constraint = keras.constraints.get(kernel_constraint)
         self._bias_constraint = keras.constraints.get(bias_constraint)
+        self._custom_build_config = {}
         self.built = False
-        # TODO: Add warning if build is implemented in subclass
-        # TODO: Add warning if call is implemented in subclass
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        subclass_build = cls.build
+
+        @functools.wraps(subclass_build)
+        def build_wrapper(self: GraphLayer, spec: tensors.GraphTensor.Spec | None):
+            GraphLayer.build(self, spec)
+            subclass_build(self, spec)
+            if not self.built and isinstance(self, keras.Model):
+                symbolic_inputs = Input(spec)
+                self.built = True
+                self(symbolic_inputs)
+
+        cls.build = build_wrapper
 
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Calls the layer.
+        """Forward pass.
 
-        Needs to be implemented by subclass.
+        Must be implemented by subclass.
 
-        Args:
+        Arguments:
             tensor:
                 A `GraphTensor` instance.
         """
-        raise NotImplementedError('`propagate` needs to be implemented.')
+        raise NotImplementedError(
+            'The forward pass of the layer is not implemented. '
+            'Please implement `propagate`.'
+        )
 
-    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
+    def build(self, tensor_spec: tensors.GraphTensor.Spec) -> None:
         """Builds the 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.
+        Optionally implemented by subclass.
 
-        Args:
-            spec:
-                A `GraphTensor.Spec` instance, corresponding to the input `GraphTensor` 
-                of the `propagate` method.
+        Arguments:
+            tensor_spec:
+                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
-
-        if not self.built:
-            # Automatically build layer or model by calling it on symbolic inputs
-            self.built = True 
-            symbolic_inputs = Input(spec)
-            self(symbolic_inputs)
-    
-    def get_build_config(self) -> dict:
-        if not hasattr(self, '_custom_build_config'):
-            return super().get_build_config()
-        return self._custom_build_config
-    
-    def build_from_config(self, config: dict) -> None:
-        use_custom_build_from_config = ('spec' in config)
-        if not use_custom_build_from_config:
-            super().build_from_config(config)
-        else:
-            spec = _deserialize_spec(config['spec'])
-            self.build(spec)
+        if isinstance(tensor_spec, tensors.GraphTensor.Spec):
+            self._custom_build_config['spec'] = _serialize_spec(tensor_spec)
 
     def call(
         self, 
@@ -127,6 +143,19 @@ class GraphLayer(keras.layers.Layer):
             outputs = tensors.from_dict(outputs)
         return outputs
 
+    def get_build_config(self) -> dict:
+        if self._custom_build_config:
+            return self._custom_build_config
+        return super().get_build_config()
+    
+    def build_from_config(self, config: dict) -> None:
+        serialized_spec = config.get('spec')
+        if serialized_spec is not None:
+            spec = _deserialize_spec(serialized_spec)
+            self.build(spec)
+        else:
+            super().build_from_config(config)
+
     def get_weight(
         self,
         shape: tf.TensorShape,
@@ -168,7 +197,7 @@ class GraphLayer(keras.layers.Layer):
             use_bias=self._use_bias,
             kernel_regularizer=self._kernel_regularizer,
             bias_regularizer=self._bias_regularizer,
-            activity_regularizer=self.activity_regularizer,
+            activity_regularizer=self._activity_regularizer,
             kernel_constraint=self._kernel_constraint,
             bias_constraint=self._bias_constraint,
         )
@@ -194,52 +223,137 @@ class GraphLayer(keras.layers.Layer):
                 keras.regularizers.serialize(self._kernel_regularizer),
             "bias_regularizer": 
                 keras.regularizers.serialize(self._bias_regularizer),
+            "activity_regularizer":
+                keras.regularizers.serialize(self._activity_regularizer),
             "kernel_constraint": 
                 keras.constraints.serialize(self._kernel_constraint),
             "bias_constraint": 
                 keras.constraints.serialize(self._bias_constraint),
         })
         return config
-    
+
 
 @keras.saving.register_keras_serializable(package='molcraft')
 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`.
+    This layer implements the three basic steps of a graph neural network layer, each of which 
+    can (optionally) be overridden by the `GraphConv` subclass:
+
+    1. **message(graph)**, which computes the *messages* to be passed to target nodes;
+    2. **aggregate(graph)**, which *aggregates* messages to target nodes;
+    3. **update(graph)**, which further *updates* (target) nodes.
+
+    Note: for skip connection to work, the `GraphConv` subclass requires final node feature 
+    output dimension to be equal to `units`. 
+
+    Arguments:
+        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`.
+        use_bias (bool):
+            Whether bias should be used in dense layers. Default to `True`.
+        normalization (bool, str):
+            Whether `LayerNormalization` should be applied to the final node feature output.
+            To use `BatchNormalization`, specify `batch_norm`. Default to `False`.
+        skip_connection (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`.
+        kernel_initializer (keras.initializers.Initializer, str):
+            Initializer for the kernel weight matrix of the dense layers.
+            Default to `glorot_uniform`.
+        bias_initializer (keras.initializers.Initializer, str):
+            Initializer for the bias weight vector of the dense layers.
+            Default to `zeros`.
+        kernel_regularizer (keras.regularizers.Regularizer, None):
+            Regularizer function applied to the kernel weight matrix.
+            Default to `None`.
+        bias_regularizer (keras.regularizers.Regularizer, None):
+            Regularizer function applied to the bias weight vector.
+            Default to `None`.
+        activity_regularizer (keras.regularizers.Regularizer, None):
+            Regularizer function applied to the output of the dense layers.
+            Default to `None`.
+        kernel_constraint (keras.constraints.Constraint, None):
+            Constraint function applied to the kernel weight matrix.
+            Default to `None`.
+        bias_constraint (keras.constraints.Constraint, None):
+            Constraint function applied to the bias weight vector.
+            Default to `None`.
     """
         
     def __init__(
         self, 
-        units: int, 
-        normalize: bool = False,
-        skip_connection: bool = False, 
+        units: int = None, 
+        activation: str | keras.layers.Activation | None = None,
+        use_bias: bool = True,
+        normalization: bool | str = False,
+        skip_connection: bool | str = True, 
         **kwargs
     ) -> None:
-        super().__init__(**kwargs)
-        self.units = units
-        self._normalize_aggregate = normalize
+        super().__init__(use_bias=use_bias, **kwargs)
+        self._units = units
+        self._normalization = normalization
         self._skip_connection = skip_connection
+        self._activation = keras.activations.get(activation)
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        subclass_build = cls.build
+
+        @functools.wraps(subclass_build)
+        def build_wrapper(self, spec):
+            GraphConv.build(self, spec)
+            return subclass_build(self, spec)
+
+        cls.build = build_wrapper
+
+    @property 
+    def units(self):
+        return self._units 
+    
+    def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        """Forward pass.
+
+        Invokes `message(graph)`, `aggregate(graph)` and `update(graph)` in sequence.
+
+        Arguments:
+            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)
+
+        updated_node_feature = tensor.node['feature']
+
+        if self._skip_connection:
+            if self._use_weighted_skip_connection:
+                input_node_feature *= self._skip_connection_weight
+            updated_node_feature += input_node_feature
+        
+        if self._normalization:
+            updated_node_feature = self._output_norm(updated_node_feature)
+
+        return tensor.update({'node': {'feature': updated_node_feature}})
     
     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._skip_connection and (node_feature_dim != self.units)
@@ -254,81 +368,115 @@ class GraphConv(GraphLayer):
             self._residual_projection = self.get_dense(
                 self.units, name='residual_projection'
             )
-        if self._normalize_aggregate:
-            self._aggregation_norm = keras.layers.LayerNormalization(
-                name='aggregation_normalizer'
+
+        skip_connection = str(self._skip_connection).lower()
+        self._use_weighted_skip_connection = skip_connection.startswith('weight')
+        if self._use_weighted_skip_connection:
+            self._skip_connection_weight = self.add_weight(
+                name='skip_connection_weight', 
+                shape=(), 
+                initializer='ones',
+                trainable=True,
             )
-            self._aggregation_norm.build([None, self.units])
 
-        super().build(spec)
+        if self._normalization:
+            if str(self._normalization).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 = 'edge' in spec.edge 
+
+        has_overridden_message = self.__class__.message != GraphConv.message 
+        if not has_overridden_message:
+            self._message_dense = self.get_dense(self.units)
+
+        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
 
-    @abc.abstractmethod
     def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         """Compute messages.
 
-        This method needs to be implemented by subclass.
+        This method may be overridden by subclass.
 
-        Args:
+        Arguments:
             tensor:
                 The inputted `GraphTensor` instance.
         """
-    
-    @abc.abstractmethod
+        if not self._has_edge_feature:
+            message_feature = tensor.gather('feature', 'source')
+        else:
+            message_feature = keras.ops.concatenate(
+                [
+                    tensor.gather('feature', 'source'),
+                    tensor.edge['feature']
+                ], 
+                axis=-1
+            )
+        message = self._message_dense(message_feature)
+        return tensor.update(
+            {
+                'edge': {
+                    'message': message
+                }
+            }
+        )
+
     def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         """Aggregates messages.
 
-        This method needs to be implemented by subclass.
+        This method may be overridden by subclass.
 
-        Args:
+        Arguments:
             tensor:
                 A `GraphTensor` instance containing a message.
         """
+        aggregate = tensor.aggregate('message', mode='mean')
+        return tensor.update(
+            {
+                'node': {
+                    'feature': aggregate,
+                    'previous_feature': tensor.node['feature']
+                },
+                'edge': {
+                    'message': None
+                }
+            }
+        )
 
-    @abc.abstractmethod
     def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         """Updates nodes. 
 
-        This method needs to be implemented by subclass.
+        This method may be overridden by subclass.
 
-        Args:
+        Arguments:
             tensor:
                 A `GraphTensor` instance containing aggregated messages 
                 (updated node features).
         """
-
-    def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Calls the layer.
-
-        The `GraphConv` layer invokes `message`, `aggregate` and `update`
-        in sequence.
-        
-        Args:
-            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)
-
-        if self._normalize_aggregate:
-            normalized_node_feature = self._aggregation_norm(tensor.node['feature'])
-            tensor = tensor.update({'node': {'feature': normalized_node_feature}})
-
-        tensor = self.update(tensor)
-
-        if not self._skip_connection:
-            return tensor
-        
-        updated_node_feature = tensor.node['feature']
+        if not 'previous_feature' in tensor.node:
+            feature = tensor.node['feature']
+        else:
+            feature = keras.ops.concatenate(
+                [
+                    tensor.node['feature'],
+                    tensor.node['previous_feature']
+                ],
+                axis=-1
+            )
+        update = self._output_dense(feature)
+        update = self._output_activation(update)
         return tensor.update(
             {
                 'node': {
-                    'feature': updated_node_feature + input_node_feature
+                    'feature': update,
+                    'previous_feature': None,
                 }
             }
         )
@@ -337,16 +485,44 @@ class GraphConv(GraphLayer):
         config = super().get_config()
         config.update({
             'units': self.units,
-            'normalize': self._normalize_aggregate,
+            'activation': keras.activations.serialize(self._activation),
+            'normalization': self._normalization,
             'skip_connection': self._skip_connection,
         })
         return config
     
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class GINConv(GraphConv):
+class GIConv(GraphConv):
 
     """Graph isomorphism network layer.
+
+    >>> graph = molcraft.tensors.GraphTensor(
+    ...     context={
+    ...         'size': [2]
+    ...     },
+    ...     node={
+    ...         'feature': [[1.], [2.]]
+    ...     },
+    ...     edge={
+    ...         'source': [0, 1],
+    ...         'target': [1, 0],
+    ...     }
+    ... )
+    >>> conv = molcraft.layers.GIConv(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__(
@@ -354,24 +530,20 @@ class GINConv(GraphConv):
         units: int,
         activation: keras.layers.Activation | str | None = 'relu',
         use_bias: bool = True,
-        normalize: bool = True,
-        dropout: float = 0.0,
+        normalization: bool = False,
         update_edge_feature: bool = True,
         **kwargs,
     ):
         super().__init__(
             units=units, 
-            normalize=normalize, 
+            activation=activation,
+            normalization=normalization, 
             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.
-        """
+    def build(self, spec: tensors.GraphTensor.Spec) -> None:
         node_feature_dim = spec.node['feature'].shape[-1]
 
         self.epsilon = self.add_weight(
@@ -381,7 +553,8 @@ class GINConv(GraphConv):
             trainable=True,
         )
 
-        if 'feature' in spec.edge:
+        self._has_edge_feature = 'feature' in spec.edge
+        if self._has_edge_feature:
             edge_feature_dim = spec.edge['feature'].shape[-1]
 
             if not self._update_edge_feature:
@@ -395,31 +568,21 @@ class GINConv(GraphConv):
 
             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 != GINConv.update 
+        has_overridden_update = self.__class__.update != GIConv.update 
         if not has_overridden_update:
-            # Use default feedforward network
-
-            self._feedforward_dropout = keras.layers.Dropout(self._dropout)
+            self._feedforward_intermediate_dense = self.get_dense(self.units)
             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:
-        """Computes 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:
+        if self._has_edge_feature:
             message += edge_feature
         return tensor.update(
             {
@@ -431,12 +594,8 @@ class GINConv(GraphConv):
         )
 
     def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Aggregates messages.
-        """
-        node_feature = tensor.aggregate('message')
+        node_feature = tensor.aggregate('message', mode='mean')
         node_feature += (1 + self.epsilon) * tensor.node['feature']
-        node_feature = self._feedforward_intermediate_dense(node_feature)
-        node_feature = self._feedforward_activation(node_feature)
         return tensor.update(
             {
                 'node': {
@@ -449,10 +608,9 @@ class GINConv(GraphConv):
         )
     
     def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Updates nodes. 
-        """
         node_feature = tensor.node['feature']
-        node_feature = self._feedforward_dropout(node_feature)
+        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(
             {
@@ -465,17 +623,217 @@ class GINConv(GraphConv):
     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='molgraphx')
+class GAConv(GraphConv):
+
+    """Graph attention network layer.
+
+    >>> graph = molcraft.tensors.GraphTensor(
+    ...     context={
+    ...         'size': [2]
+    ...     },
+    ...     node={
+    ...         'feature': [[1.], [2.]]
+    ...     },
+    ...     edge={
+    ...         'source': [0, 1],
+    ...         'target': [1, 0],
+    ...     }
+    ... )
+    >>> conv = molcraft.layers.GAConv(units=4, heads=2)
+    >>> 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,
+        heads: int = 8,
+        activation: keras.layers.Activation | str | None = "relu",
+        use_bias: bool = True,
+        normalization: bool = False,
+        update_edge_feature: bool = True,
+        attention_activation: keras.layers.Activation | str | None = "leaky_relu",
+        **kwargs,
+    ) -> None:
+        kwargs['skip_connection'] = False
+        super().__init__(
+            units=units, 
+            activation=activation,
+            use_bias=use_bias,
+            normalization=normalization, 
+            **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._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(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
+        if self._update_edge_feature:
+            self._edge_dense = self.get_einsum_dense(
+                'ijh,jkh->ikh', (self.head_units, self.heads)
+            )
+        self._node_dense = self.get_einsum_dense(
+            'ij,jkh->ikh', (self.head_units, self.heads)
+        )
+        self._feature_dense = self.get_einsum_dense(
+            'ij,jkh->ikh', (self.head_units, self.heads)
+        )
+        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(
+            [
+                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', mode='sum')
+        node_feature += self._node_self_dense(tensor.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 = tensor.node['feature']
+        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
+                }
+            }
+        )
+    
+    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
+    
+
 @keras.saving.register_keras_serializable(package='molcraft')
 class GTConv(GraphConv):
 
     """Graph transformer layer.
+
+    >>> graph = molcraft.tensors.GraphTensor(
+    ...     context={
+    ...         'size': [2]
+    ...     },
+    ...     node={
+    ...         'feature': [[1.], [2.]]
+    ...     },
+    ...     edge={
+    ...         'source': [0, 1],
+    ...         'target': [1, 0],
+    ...     }
+    ... )
+    >>> conv = molcraft.layers.GTConv(units=4, heads=2)
+    >>> 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__(
@@ -484,26 +842,22 @@ class GTConv(GraphConv):
         heads: int = 8,
         activation: keras.layers.Activation | str | None = "relu",
         use_bias: bool = True,
-        normalize: bool = True,
-        dropout: float = 0.0,
+        normalization: bool = False,
         attention_dropout: float = 0.0,
         **kwargs,
     ) -> None:
-        kwargs['skip_connection'] = False
         super().__init__(
             units=units, 
-            normalize=normalize, 
+            activation=activation,
             use_bias=use_bias,
+            normalization=normalization, 
             **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):
@@ -513,68 +867,41 @@ class GTConv(GraphConv):
     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])
-            
+    def build(self, spec: tensors.GraphTensor.Spec) -> None:
         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
 
-        self._add_edge_bias = not 'bias' in spec.edge
-        if self._add_edge_bias:
-            self._add_edge_bias = AddEdgeBias()
-            self._add_edge_bias.build_from_spec(spec)
+        if self._add_bias:
+            self._edge_bias = EdgeBias(biases=self.heads)
 
         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_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:
-        """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)
@@ -587,12 +914,8 @@ class GTConv(GraphConv):
 
         attention_score = keras.ops.sum(query * key, axis=1, keepdims=True)
         attention_score /= keras.ops.sqrt(float(self.head_units))
-
-        if self._add_edge_bias:
-            tensor = self._add_edge_bias(tensor)
             
-        attention_score += keras.ops.expand_dims(tensor.edge['bias'], -1)
-        
+        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)
 
@@ -606,12 +929,9 @@ class GTConv(GraphConv):
         )
 
     def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Aggregates messages.
-        """
-        node_feature = tensor.aggregate('message')
+        node_feature = tensor.aggregate('message', mode='sum')
         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': {
@@ -626,49 +946,257 @@ class GTConv(GraphConv):
         )
 
     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)
+        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({
+            "heads": self._heads,
+            'attention_dropout': self._attention_dropout,
+        })
+        return config
+    
+
+@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,
+        normalization: bool = False,
+        **kwargs
+    ) -> None:
+        super().__init__(
+            units=units, 
+            activation=activation,
+            use_bias=use_bias,
+            normalization=normalization, 
+            **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
+            )
+
+    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')
+        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({})
+        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']
 
-        residual = tensor.node['residual']
-        if self.project_residual:
-            residual = self._residual_dense(residual)
+        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)
 
-        node_feature += residual
-        residual = node_feature
+            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)
 
-        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)
+        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:
+        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': {
+                    'feature': node_feature,
+                    'residual': tensor.node['feature']
+                },
+                'edge': {
+                    'message': None,
+                    'weight': None,
+                }
+            }
+        )
+    
 
-        node_feature += residual
+@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(
             {
-                'node': {
-                    'feature': node_feature,
-                },
+                'edge': {
+                    'message': message,
+                }
             }
         )
     
-    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 EGConv3D(GraphConv):
 
-    """Equivariant graph neural network layer.
+    """Equivariant graph neural network layer 3D.
     """
 
     def __init__(
@@ -676,48 +1204,33 @@ class EGConv3D(GraphConv):
         units: int = 128, 
         activation: keras.layers.Activation | str | None = None, 
         use_bias: bool = True,
-        normalize: bool = True,
-        dropout: float = 0.0,
+        normalization: bool = False,
         **kwargs
     ) -> None:
         super().__init__(
             units=units, 
-            normalize=normalize, 
+            activation=activation,
             use_bias=use_bias,
+            normalization=normalization, 
             **kwargs
         )
-        self._activation = keras.activations.get(activation)
-        self._dropout = dropout or 0.0
 
-    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
+    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.'
             )
-        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._has_edge_feature = 'feature' in spec.edge
         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.output_dense = self.get_dense(self.units)
 
     def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Computes messages.
-        """
         relative_node_coordinate = keras.ops.subtract(
             tensor.gather('coordinate', 'target'), 
             tensor.gather('coordinate', 'source')
@@ -760,8 +1273,6 @@ class EGConv3D(GraphConv):
         )
 
     def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Aggregates messages.
-        """
         coefficient = keras.ops.bincount(
             tensor.edge['source'], 
             minlength=tensor.num_nodes
@@ -776,7 +1287,7 @@ class EGConv3D(GraphConv):
         updated_coordinate = tensor.aggregate('relative_node_coordinate') * coefficient
         updated_coordinate += tensor.node['coordinate']
 
-        aggregate = tensor.aggregate('message')
+        aggregate = tensor.aggregate('message', mode='mean')
         return tensor.update(
             {
                 'node': {
@@ -792,8 +1303,6 @@ class EGConv3D(GraphConv):
         ) 
 
     def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Updates nodes.
-        """
         updated_node_feature = self.update_fn(
             keras.ops.concatenate(
                 [
@@ -803,7 +1312,7 @@ class EGConv3D(GraphConv):
                 axis=-1
             )
         )  
-        updated_node_feature = self._dropout_layer(updated_node_feature)
+        updated_node_feature = self.output_dense(updated_node_feature)
         return tensor.update(
             {
                 'node': {
@@ -815,65 +1324,46 @@ class EGConv3D(GraphConv):
     
     def get_config(self) -> dict:
         config = super().get_config()
-        config.update({
-            'activation': keras.activations.serialize(self._activation),
-            'dropout': self._dropout,
-        })
+        config.update({})
         return config 
     
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class Projection(GraphLayer):
-    """Base graph projection layer.
+class Readout(GraphLayer):
+
+    """Readout 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])
+    def __init__(self, mode: str | None = None, **kwargs):
+        kwargs['kernel_initializer'] = None 
+        kwargs['bias_initializer'] = None
+        super().__init__(**kwargs)
+        self.mode = mode
+        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 
+        elif str(self.mode).lower().startswith('super'):
+            self._reduce_fn = keras.ops.segment_sum
+        else:
+            self._reduce_fn = ops.segment_mean
 
-    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 propagate(self, tensor: tensors.GraphTensor) -> tf.Tensor:
+        node_feature = tensor.node['feature']
+        if str(self.mode).lower().startswith('super'):
+            node_feature = keras.ops.where(
+                tensor.node['super'][:, None], node_feature, 0.0
+            )
+        return self._reduce_fn(
+            node_feature, tensor.graph_indicator, tensor.num_subgraphs
+        )
 
     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
+        config['mode'] = self.mode 
+        return config 
+    
     
-
 @keras.saving.register_keras_serializable(package='molcraft')
 class GraphNetwork(GraphLayer):
 
@@ -881,7 +1371,7 @@ class GraphNetwork(GraphLayer):
 
     Sequentially calls graph layers (`GraphLayer`) and concatenates its output. 
 
-    Args:
+    Arguments:
         layers (list):
             A list of graph layers.
     """
@@ -891,36 +1381,32 @@ class GraphNetwork(GraphLayer):
         self.layers = layers
         self._update_edge_feature = False
 
-    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
-        """Builds the layer.
-        """
+    def build(self, spec: tensors.GraphTensor.Spec) -> None:
         units = self.layers[0].units 
         node_feature_dim = spec.node['feature'].shape[-1]
-        if node_feature_dim != units:
+        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._update_node_feature = True 
-        has_edge_feature = 'feature' in spec.edge 
-        if has_edge_feature:
+        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:
+            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)
-                self._update_edge_feature = 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:
+        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:
@@ -945,7 +1431,7 @@ class GraphNetwork(GraphLayer):
         Performs the same forward pass as `propagate` but with a `GradientTape`
         watching intermediate node features.
 
-        Args:
+        Arguments:
             tensor (tensors.GraphTensor):
                 The graph input.
         """
@@ -1003,24 +1489,25 @@ class NodeEmbedding(GraphLayer):
     def __init__(
         self, 
         dim: int = None, 
+        normalization: bool = False,
         embed_context: bool = True,
+        allow_reconstruction: bool = False,
         allow_masking: bool = True, 
         **kwargs
     ) -> None:
         super().__init__(**kwargs)
         self.dim = dim
+        self._normalization = normalization
         self._embed_context = embed_context
         self._masking_rate = None
         self._allow_masking = allow_masking
+        self._allow_reconstruction = allow_reconstruction
 
-    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
-        """Builds the layer.
-        """
+    def build(self, spec: tensors.GraphTensor.Spec) -> None:
         feature_dim = spec.node['feature'].shape[-1]
         if not self.dim:
             self.dim = feature_dim
         self._node_dense = self.get_dense(self.dim)
-        self._node_dense.build([None, feature_dim])
 
         self._has_super = 'super' in spec.node
         has_context_feature = 'feature' in spec.context
@@ -1034,11 +1521,18 @@ class NodeEmbedding(GraphLayer):
         if self._embed_context:
             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._normalization:
+            if str(self._normalization).lower().startswith('batch'):
+                self._norm = keras.layers.BatchNormalization(
+                    name='output_batch_norm'
+                )
+            else:
+                self._norm = keras.layers.LayerNormalization(
+                    name='output_layer_norm'
+                )
 
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Calls the layer.
-        """
         feature = self._node_dense(tensor.node['feature'])
 
         if self._has_super:
@@ -1068,8 +1562,13 @@ class NodeEmbedding(GraphLayer):
             # Slience warning of 'no gradients for variables'
             feature = feature + (self._mask_feature * 0.0)
 
-        return tensor.update({'node': {'feature': feature}})
+        if self._normalization:
+            feature = self._norm(feature)
 
+        if not self._allow_reconstruction:
+            return tensor.update({'node': {'feature': feature}})
+        return tensor.update({'node': {'feature': feature, 'target_feature': feature}})
+    
     @property 
     def masking_rate(self):
         return self._masking_rate 
@@ -1087,7 +1586,10 @@ class NodeEmbedding(GraphLayer):
         config = super().get_config()
         config.update({
             'dim': self.dim,
-            'allow_masking': self._allow_masking
+            'normalization': self._normalization,
+            'embed_context': self._embed_context,
+            'allow_masking': self._allow_masking,
+            'allow_reconstruction': self._allow_reconstruction,
         })
         return config
     
@@ -1103,22 +1605,21 @@ class EdgeEmbedding(GraphLayer):
     def __init__(
         self, 
         dim: int = None, 
+        normalization: bool = False,
         allow_masking: bool = True, 
         **kwargs
     ) -> None:
         super().__init__(**kwargs)
         self.dim = dim
+        self._normalization = normalization
         self._masking_rate = None
         self._allow_masking = allow_masking
 
-    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
-        """Builds the layer.
-        """
+    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.build([None, feature_dim])
 
         self._has_super = 'super' in spec.edge
         if self._has_super:
@@ -1126,9 +1627,17 @@ class EdgeEmbedding(GraphLayer):
         if self._allow_masking:
             self._mask_feature = self.get_weight(shape=[self.dim], name='mask_edge_feature')
 
+        if self._normalization:
+            if str(self._normalization).lower().startswith('batch'):
+                self._norm = keras.layers.BatchNormalization(
+                    name='output_batch_norm'
+                )
+            else:
+                self._norm = keras.layers.LayerNormalization(
+                    name='output_layer_norm'
+                )
+
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        """Calls the layer.
-        """
         feature = self._edge_dense(tensor.edge['feature'])
 
         if self._has_super:
@@ -1153,7 +1662,10 @@ class EdgeEmbedding(GraphLayer):
             # Slience warning of 'no gradients for variables'
             feature = feature + (self._mask_feature * 0.0)
 
-        return tensor.update({'edge': {'feature': feature}})
+        if self._normalization:
+            feature = self._norm(feature)
+
+        return tensor.update({'edge': {'feature': feature, 'embedding': feature}})
 
     @property 
     def masking_rate(self):
@@ -1172,17 +1684,67 @@ class EdgeEmbedding(GraphLayer):
         config = super().get_config()
         config.update({
             'dim': self.dim,
+            'normalization': self._normalization,
             '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):
-        super().__init__(units=units, activation=activation, field='context', **kwargs)
+        kwargs['field'] = 'context'
+        super().__init__(units=units, activation=activation, **kwargs)
 
 
 @keras.saving.register_keras_serializable(package='molcraft')
@@ -1190,7 +1752,8 @@ 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)
+        kwargs['field'] = 'node'
+        super().__init__(units=units, activation=activation, **kwargs)
 
 
 @keras.saving.register_keras_serializable(package='molcraft')
@@ -1198,103 +1761,126 @@ 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)
-    
+        kwargs['field'] = 'edge'
+        super().__init__(units=units, activation=activation, **kwargs)
+
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class Readout(keras.layers.Layer):
+class Reconstruction(GraphLayer):
 
-    def __init__(self, mode: str | None = None, **kwargs):
+    def __init__(
+        self, 
+        loss: keras.losses.Loss | str = 'mse',
+        loss_weight: float = 0.5, 
+        **kwargs
+    ):
         super().__init__(**kwargs)
-        self.mode = mode
-        if not self.mode:
-            self._reduce_fn = None
-        elif 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 
-        elif str(self.mode).lower().startswith('super'):
-            self._reduce_fn = keras.ops.segment_sum
-        else:
-            self._reduce_fn = ops.segment_mean
-
-    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
-        """Builds the layer.
-        """
-        pass
-
-    def reduce(self, tensor: tensors.GraphTensor) -> tf.Tensor:
-        if self._reduce_fn is None:
-            raise NotImplementedError("Need to define a reduce method.")
-        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
+        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`.'
             )
-        return self._reduce_fn(
-            tensor.node['feature'], tensor.graph_indicator, tensor.num_subgraphs
-        )
+         output_dim = spec.node['target_feature'].shape[-1]
+         self._dense = self.get_dense(output_dim)
 
-    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 propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        target_node_feature = tensor.node['target_feature']
+        transformed_node_feature = tensor.node['feature']
 
-    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)
+        reconstructed_node_feature = self._dense(
+            transformed_node_feature
+        )
 
-    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
+        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) -> dict:
+    def get_config(self):
         config = super().get_config()
-        config['mode'] = self.mode 
-        return 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 AddEdgeBias(GraphLayer):
+class EdgeBias(GraphLayer):
 
-    def build_from_spec(self, spec: tensors.GraphTensor.Spec) -> None:
+    def __init__(self, biases: int, **kwargs):
+        super().__init__(**kwargs)
+        self.biases = biases
+
+    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(units=1)
+            self._edge_feature_dense = self.get_dense(self.biases)
         if self._has_edge_length:
             self._edge_length_dense = self.get_dense(
-                units=1, kernel_initializer='zeros'
+                self.biases, kernel_initializer='zeros'
             )
-            
+
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         bias = keras.ops.zeros(
-            shape=(tensor.num_edges, 1), 
+            shape=(tensor.num_edges, self.biases), 
             dtype=tensor.node['feature'].dtype
         )
         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 tensor.update(
-            {
-                'edge': {
-                    'bias': bias
-                }
-            }
+        return bias
+
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({'biases': self.biases})
+        return config
+    
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class GaussianDistance(GraphLayer):
+
+    def __init__(self, kernels: int, **kwargs):
+        super().__init__(**kwargs)
+        self.kernels = kernels 
+
+    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
+        )
+
+    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({
+            'kernels': self.kernels,
+        })
+        return config
     
 
 def Input(spec: tensors.GraphTensor.Spec) -> dict:
@@ -1412,13 +1998,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