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 value p 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)

reset_parameters()[source]

Resets all learnable parameters of the module.

forward(x: Tensor, index: = None, ptr: = None, dim_size: = None, dim: int = -2) [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 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)