torch_geometric.nn.aggr.MLPAggregation

class MLPAggregation(in_channels: int, out_channels: int, max_num_elements: int, **kwargs)[source]

Bases: Aggregation

Performs MLP aggregation in which the elements to aggregate are flattened into a single vectorial representation, and are then processed by a Multi-Layer Perceptron (MLP), as described in the “Graph Neural Networks with Adaptive Readouts” paper.

Note

MLPAggregation requires sorted indices index as input. Specifically, if you use this aggregation as part of MessagePassing, ensure that edge_index is sorted by destination nodes, either by manually sorting edge indices via sort_edge_index() or by calling torch_geometric.data.Data.sort().

Warning

MLPAggregation is not a permutation-invariant operator.

Parameters:

  • in_channels (int) – Size of each input sample.

  • out_channels (int) – Size of each output sample.

  • max_num_elements (int) – The maximum number of elements to aggregate per group.

  • **kwargs (optional) – Additional arguments of torch_geometric.nn.models.MLP.

reset_parameters()[source]

Resets all learnable parameters of the module.

forward(x: Tensor, index: Optional[Tensor] = None, ptr: Optional[Tensor] = None, dim_size: Optional[int] = None, dim: int = -2) → Tensor[source]

Forward pass.

Parameters:

  • x (torch.Tensor) – The source tensor.

  • index (torch.Tensor, optional) – The indices of elements for applying the aggregation. One of index or ptr must be defined. (default: None)

  • ptr (torch.Tensor, optional) – If given, computes the aggregation based on sorted inputs in CSR representation. One of index or ptr must be defined. (default: None)

  • dim_size (int, optional) – The size of the output tensor at dimension dim after aggregation. (default: None)

  • dim (int, optional) – The dimension in which to aggregate. (default: -2)

  • max_num_elements – (int, optional): The maximum number of elements within a single aggregation group. (default: None)