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.

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]
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)