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.
docstr: Prepare gnn sublayer stack.
  
  :param in_dim: Input dimensionality of layer.
  :param out_dim: Output dimensionality of layer.
  :return: Gnn sublayer stack.
docstr_tok: ["Prepare","gnn","sublayer","stack",".",":","param","in_dim",":","Input","dimensionality","of","layer",".",":","param","out_dim",":","Output","dimensionality","of","layer",".",":","return",":","Gnn","sublayer","stack","."]
code: def _make_sub_layer(self, in_dim: int, out_dim: int) -> nn.Sequential:
      """Prepare gnn sublayer stack.
  
      :param in_dim: Input dimensionality of layer.
      :param out_dim: Output dimensionality of layer.
      :return: Gnn sublayer stack.
      """
      sub_layers = [nn.Linear(in_dim, out_dim)]
      if self.with_layer_norm:
          sub_layers.append(nn.LayerNorm(out_dim))
      sub_layers.append(self.non_lin())
      return nn.Sequential(*sub_layers)
  
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 5 for calling the GNNBlock obj's underlying member method code with example usage: obj._make_sub_layer(in_dim, out_dim) and returns: nn