torch_geometric.nn.aggr.PatchTransformerAggregation

class PatchTransformerAggregation(in_channels: int, out_channels: int, patch_size: int, hidden_channels: int, num_transformer_blocks: int = 1, heads: int = 1, dropout: float = 0.0, aggr: Union[str, List[str]] = 'mean')[source]

Bases: Aggregation

Performs patch transformer aggregation in which the elements to aggregate are processed by multi-head attention blocks across patches, as described in the “Simplifying Temporal Heterogeneous Network for Continuous-Time Link Prediction” paper.

Parameters:

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

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

  • patch_size (int) – Number of elements in a patch.

  • hidden_channels (int) – Intermediate size of each sample.

  • num_transformer_blocks (int, optional) – Number of transformer blocks (default: 1).

  • heads (int, optional) – Number of multi-head-attentions. (default: 1)

  • dropout (float, optional) – Dropout probability of attention weights. (default: 0.0)

  • aggr (str or list[str], optional) – The aggregation module, e.g., "sum", "mean", "min", "max", "var", "std". (default: "mean")

reset_parameters() → None[source]

Resets all learnable parameters of the module.

Return type:

None

forward(x: Tensor, index: Tensor, 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 (Optional[int], default: None) – (int, optional): The maximum number of elements within a single aggregation group. (default: None)

Return type:

Tensor