torch_geometric.nn.aggr.AttentionalAggregation
- class AttentionalAggregation(gate_nn: Module, nn: Optional[Module] = None)[source]
Bases:
Aggregation
The soft attention aggregation layer from the “Graph Matching Networks for Learning the Similarity of Graph Structured Objects” paper.
\[\mathbf{r}_i = \sum_{n=1}^{N_i} \mathrm{softmax} \left( h_{\mathrm{gate}} ( \mathbf{x}_n ) \right) \cdot h_{\mathbf{\Theta}} ( \mathbf{x}_n ),\]where \(h_{\mathrm{gate}} \colon \mathbb{R}^F \to \mathbb{R}\) and \(h_{\mathbf{\Theta}}\) denote neural networks, i.e. MLPs.
- Parameters:
gate_nn (torch.nn.Module) – A neural network \(h_{\mathrm{gate}}\) that computes attention scores by mapping node features
x
of shape[-1, in_channels]
to shape[-1, 1]
(for node-level gating) or[1, out_channels]
(for feature-level gating), e.g., defined bytorch.nn.Sequential
.nn (torch.nn.Module, optional) – A neural network \(h_{\mathbf{\Theta}}\) that maps node features
x
of shape[-1, in_channels]
to shape[-1, out_channels]
before combining them with the attention scores, e.g., defined bytorch.nn.Sequential
. (default:None
)
- 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: