torch_geometric.nn.aggr.Set2Set

class Set2Set(in_channels: int, processing_steps: int, **kwargs)

Bases: Aggregation

The Set2Set aggregation operator based on iterative content-based attention, as described in the “Order Matters: Sequence to sequence for Sets” paper.

\[ \begin{align}\begin{aligned}\mathbf{q}_t &= \mathrm{LSTM}(\mathbf{q}^{*}_{t-1})\\\alpha_{i,t} &= \mathrm{softmax}(\mathbf{x}_i \cdot \mathbf{q}_t)\\\mathbf{r}_t &= \sum_{i=1}^N \alpha_{i,t} \mathbf{x}_i\\\mathbf{q}^{*}_t &= \mathbf{q}_t \, \Vert \, \mathbf{r}_t,\end{aligned}\end{align} \]

where \(\mathbf{q}^{*}_T\) defines the output of the layer with twice the dimensionality as the input.

Parameters:

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

• processing_steps (int) – Number of iterations \(T\).

• **kwargs (optional) – Additional arguments of torch.nn.LSTM.

reset_parameters()

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

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)