torch_geometric.nn.aggr.PowerMeanAggregation
- class PowerMeanAggregation(p: float = 1.0, learn: bool = False, channels: int = 1)[source]
Bases:
Aggregation
The powermean aggregation operator based on a power term, as described in the “DeeperGCN: All You Need to Train Deeper GCNs” paper.
\[\mathrm{powermean}(\mathcal{X}|p) = \left(\frac{1}{|\mathcal{X}|} \sum_{\mathbf{x}_i\in\mathcal{X}}\mathbf{x}_i^{p}\right)^{1/p},\]where \(p\) controls the power of the powermean when aggregating over a set of features \(\mathcal{X}\).
- Parameters:
p (float, optional) – Initial power for powermean aggregation. (default:
1.0
)learn (bool, optional) – If set to
True
, will learn the valuep
for powermean aggregation dynamically. (default:False
)channels (int, optional) – Number of channels to learn from \(p\). If set to a value greater than
1
, \(p\) will be learned per input feature channel. This requires compatible shapes for the input to the forward calculation. (default:1
)
- 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
)