Fix to allow use edges for MPNN layer

Author: s-maddrellmander

expts/hydra-configs/model/mpnn.yaml

@@ -22,3 +22,5 @@ architecture:
       attn_type: "none" # "full-attention", "none"
       # biased_attention: false
       attn_kwargs: null
+    virtual_node: 'sum'
+    use_virtual_edges: true

graphium/nn/architectures/global_architectures.py

@@ -12,6 +12,7 @@
 from torch import Tensor, nn
 import torch
 from torch_geometric.data import Data
+from omegaconf import DictConfig, OmegaConf
 
 # graphium imports
 from graphium.data.utils import get_keys
@@ -592,6 +593,26 @@ def _check_bad_arguments(self):
             (self.in_dim_edges > 0) or (self.full_dims_edges is not None)
         ) and not self.layer_class.layer_supports_edges:
             raise ValueError(f"Cannot use edge features with class `{self.layer_class}`")
+    
+    def get_nested_key(self, d, target_key):
+        """
+        Get the value associated with a key in a nested dictionary.
+        
+        Parameters:
+        - d: The dictionary to search in
+        - target_key: The key to search for
+        
+        Returns:
+        - The value associated with the key if found, None otherwise
+        """
+        if target_key in d:
+            return d[target_key]
+        for key, value in d.items():
+            if isinstance(value, (dict, DictConfig)):
+                nested_result = self.get_nested_key(value, target_key)
+                if nested_result is not None:
+                    return nested_result
+        return None
 
     def _create_layers(self):
         r"""
@@ -632,15 +653,18 @@ def _create_layers(self):
 
             # Find the edge key-word arguments depending on the layer type and residual connection
             this_edge_kwargs = {}
+            # import ipdb; ipdb.set_trace()
             if self.layer_class.layer_supports_edges and self.in_dim_edges > 0:
                 this_edge_kwargs["in_dim_edges"] = this_in_dim_edges
                 if "out_dim_edges" in inspect.signature(self.layer_class.__init__).parameters.keys():
                     if self.full_dims_edges is not None:
                         this_out_dim_edges = self.full_dims_edges[ii + 1]
                         this_edge_kwargs["out_dim_edges"] = this_out_dim_edges
                     else:
-                        this_out_dim_edges = self.layer_kwargs.get("out_dim_edges")
+                        this_out_dim_edges = self.get_nested_key(self.layer_kwargs, "out_dim_edges")
+                        this_edge_kwargs["out_dim_edges"] = this_out_dim_edges
                     layer_out_dims_edges.append(this_out_dim_edges)
+            # import ipdb; ipdb.set_trace()
 
             # Create the GNN layer
             self.layers.append(
@@ -659,6 +683,7 @@ def _create_layers(self):
 
             # Create the Virtual Node layer, except at the last layer
             if ii < len(residual_out_dims):
+                # import ipdb; ipdb.set_trace()
                 self.virtual_node_layers.append(
                     self.virtual_node_class(
                         in_dim=this_out_dim * self.layers[-1].out_dim_factor,