torch_geometric.nn.aggr.GRUAggregation

class GRUAggregation(in_channels: int, out_channels: int, **kwargs)[source]

Bases: Aggregation

Performs GRU aggregation in which the elements to aggregate are interpreted as a sequence, as described in the “Graph Neural Networks with Adaptive Readouts” paper.

Note

MLPAggregation requires sorted indices index as input. Specifically, if you use this aggregation as part of MessagePassing, ensure that edge_index is sorted by destination nodes, either by manually sorting edge indices via sort_edge_index() or by calling torch_geometric.data.Data.sort().

Warning

GRUAggregation is not a permutation-invariant operator.

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

  • out_channels (int) – Size of each output sample.

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

reset_parameters()[source]

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, max_num_elements: Optional[int] = None) 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 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)