GNNBlock¶
-
class
maze.perception.blocks.feed_forward.graph_nn.
GNNBlock
(*args: Any, **kwargs: Any)¶ A customizable graph neural network (GNN) block.
- Parameters
in_keys – One key identifying the input tensors.
out_keys – One key identifying the output tensors.
in_shapes – List of input shapes.
edges – List of graph edges required for message passing (aggregation).
aggregate – The aggregation function to use (max, mean, sum).
hidden_units – List containing the number of hidden units for hidden layers.
non_lin – The non-linearity to apply after each layer.
with_layer_norm – If True layer normalization is applied.
node2node_aggr – If True node to node message passing is applied.
edge2node_aggr – If True edge to node message passing is applied.
node2edge_aggr – If True node to edge message passing is applied.
edge2edge_aggr – If True edge to edge message passing is applied.
with_node_embedding – If True the node embedding is computed.
with_edge_embedding – If True the edge embedding is computed.
-
normalized_forward
(block_input: Dict[str, torch.Tensor]) → Dict[str, torch.Tensor]¶ (overrides
ShapeNormalizationBlock
)implementation of
ShapeNormalizationBlock
interface