torch_geometric.nn.aggr.VarAggregation
- class VarAggregation(semi_grad: bool = False)[source]
Bases:
Aggregation
An aggregation operator that takes the feature-wise variance across a set of elements.
\[\mathrm{var}(\mathcal{X}) = \mathrm{mean}(\{ \mathbf{x}_i^2 : x \in \mathcal{X} \}) - \mathrm{mean}(\mathcal{X})^2.\]- Parameters:
semi_grad (bool, optional) – If set to
True
, will turn off gradient calculation during \(E[X^2]\) computation. Therefore, only semi-gradients are used during backpropagation. Useful for saving memory and accelerating backward computation. (default:False
)
- 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: