class_docstr: A customizable graph neural network (GNN) block.
  
  :param in_keys: One key identifying the input tensors.
  :param out_keys: One key identifying the output tensors.
  :param in_shapes: List of input shapes.
  :param edges: List of graph edges required for message passing (aggregation).
  :param aggregate: The aggregation function to use (max, mean, sum).
  :param hidden_units: List containing the number of hidden units for hidden layers.
  :param non_lin: The non-linearity to apply after each layer.
  :param with_layer_norm: If True layer normalization is applied.
  :param node2node_aggr: If True node to node message passing is applied.
  :param edge2node_aggr: If True edge to node message passing is applied.
  :param node2edge_aggr: If True node to edge message passing is applied.
  :param edge2edge_aggr: If True edge to edge message passing is applied.
  :param with_node_embedding: If True the node embedding is computed.
  :param with_edge_embedding: If True the edge embedding is computed.
code: def __init__(
      self,
      in_keys: Union[str, List[str]],
      out_keys: Union[str, List[str]],
      in_shapes: Union[Sequence[int], List[Sequence[int]]],
      edges: List[Tuple[int, int]],
      aggregate: str,
      hidden_units: List[int],
      non_lin: Union[str, type(nn.Module)],
      with_layer_norm: bool,
      node2node_aggr: bool,
      edge2node_aggr: bool,
      node2edge_aggr: bool,
      edge2edge_aggr: bool,
      with_node_embedding: bool,
      with_edge_embedding: bool,
  ):
      super().__init__(
          in_keys=in_keys,
          out_keys=out_keys,
          in_shapes=in_shapes,
          in_num_dims=[3, 3],
          out_num_dims=[3, 3],
      )
  
      edge_list = np.vstack(edges).T
  
      self.num_nodes = in_shapes[0][0]
      self.num_node_features = in_shapes[0][1]
      self.num_edge_features = in_shapes[1][1]
  
      self.num_edges = edge_list.shape[1]
      self.edge_list = edge_list
  
      self.node2node_aggr = node2node_aggr
      self.edge2node_aggr = edge2node_aggr
      self.node2edge_aggr = node2edge_aggr
      self.edge2edge_aggr = edge2edge_aggr
  
      self.with_node_embedding = with_node_embedding
      self.with_edge_embedding = with_edge_embedding
  
      self.hidden_units = hidden_units
      self.n_layers = len(self.hidden_units)
      self.non_lin = Factory(base_type=nn.Module).type_from_name(
          non_lin
      )
      self.with_layer_norm = with_layer_norm
  
      # sort edges
      self.edge_list = np.sort(self.edge_list, axis=0)
  
      # node layers
      if with_node_embedding or node2edge_aggr or edge2node_aggr:
          self.node_layers = nn.ModuleList()
          for i in range(self.n_layers):
              if i == 0:
                  in_dim = self.num_node_features
              else:
                  in_dim = self.hidden_units[i - 1]
                  if self.node2node_aggr:
                      in_dim += self.hidden_units[i - 1]
                  if self.edge2node_aggr:
                      in_dim += self.hidden_units[i - 1]
  
              if with_node_embedding or (i + 1) < self.n_layers:
                  self.node_layers.append(
                      self._make_sub_layer(
                          in_dim=in_dim, out_dim=self.hidden_units[i]
                      )
                  )
  
      # edge layers
      if with_edge_embedding or edge2node_aggr or node2edge_aggr:
          self.edge_layers = nn.ModuleList()
          for i in range(self.n_layers):
              if i == 0:
                  in_dim = self.num_edge_features
              else:
                  in_dim = self.hidden_units[i - 1]
                  if self.edge2edge_aggr:
                      in_dim += self.hidden_units[i - 1]
                  if self.node2edge_aggr:
                      in_dim += self.hidden_units[i - 1]
  
              if with_edge_embedding or (i + 1) < self.n_layers:
                  self.edge_layers.append(
                      self._make_sub_layer(
                          in_dim=in_dim, out_dim=self.hidden_units[i]
                      )
                  )
  
      # prepare aggregation layers
      self.aggregate = aggregate
      pooling_matrices = self._prepare_pooling_matrices()
      self.n2n_aggregation = AggregationLayer(
          aggregate=aggregate, pooling_mask=pooling_matrices[0]
      )
      self.e2n_aggregation = AggregationLayer(
          aggregate=aggregate, pooling_mask=pooling_matrices[1]
      )
      self.n2e_aggregation = AggregationLayer(
          aggregate=aggregate, pooling_mask=pooling_matrices[2]
      )
      self.e2e_aggregation = AggregationLayer(
          aggregate=aggregate, pooling_mask=pooling_matrices[3]
      )
  
imports: from typing import Any, List, Tuple, Union, Sequence, Dict
  import numpy
  import torch
  from torch import nn
  from torch_scatter import scatter
  from maze.core.annotations import override
  from maze.core.utils.factory import Factory
  from maze.perception.blocks.shape_normalization import ShapeNormalizationBlock
label_desc: Use image node_id 1 to create a new GNNBlock object from inherited base classes: ShapeNormalizationBlock with example: obj = GNNBlock(in_keys, out_keys, in_shapes, edges, aggregate, hidden_units, non_lin, with_layer_norm, node2node_aggr, edge2node_aggr, node2edge_aggr, edge2edge_aggr, with_node_embedding, with_edge_embedding)