torch_geometric.nn.models.HeteroJumpingKnowledge

class HeteroJumpingKnowledge(types: List[str], mode: str, channels: Optional[int] = None, num_layers: Optional[int] = None)[source]

Bases: Module

A heterogeneous version of the JumpingKnowledge module.

Parameters:

  • types (List[str]) – The keys of the input dictionary.

  • mode (str) – The aggregation scheme to use ("cat", "max" or "lstm").

  • channels (int, optional) – The number of channels per representation. Needs to be only set for LSTM-style aggregation. (default: None)

  • num_layers (int, optional) – The number of layers to aggregate. Needs to be only set for LSTM-style aggregation. (default: None)

forward(xs_dict: Dict[str, List[Tensor]]) → Dict[str, Tensor][source]

Forward pass.

Parameters:

xs_dict (Dict[str, List[torch.Tensor]]) – A dictionary holding a list of layer-wise representation for each type.

Return type:

Dict[str, Tensor]

reset_parameters() → None[source]

Resets all learnable parameters of the module.

Return type:

None