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 indicesindex
as input. Specifically, if you use this aggregation as part ofMessagePassing
, ensure thatedge_index
is sorted by destination nodes, either by manually sorting edge indices viasort_edge_index()
or by callingtorch_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
.
- 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
orptr
must be defined. (default:None
)ptr (torch.Tensor, optional) – If given, computes the aggregation based on sorted inputs in CSR representation. One of
index
orptr
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
)
- Return type: