torch_geometric.nn.models.HeteroJumpingKnowledge

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

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]

Forward pass.

Parameters:

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

reset_parameters() → None

Resets all learnable parameters of the module.