SimpleMessagePassingRepresentation

class SimpleMessagePassingRepresentation(triples_factory: CoreTriplesFactory, layers: str | None | type[None] | Sequence[str | None | type[None]], layers_kwargs: Mapping[str, Any] | None | Sequence[Mapping[str, Any] | None] = None, base: str | Representation | type[Representation] | None = None, base_kwargs: Mapping[str, Any] | None = None, max_id: int | None = None, shape: int | Sequence[int] | None = None, activations: str | Module | type[Module] | None | Sequence[str | Module | type[Module] | None] = None, activations_kwargs: Mapping[str, Any] | None | Sequence[Mapping[str, Any] | None] = None, restrict_k_hop: bool = False, **kwargs)[source]

Bases: MessagePassingRepresentation

A representation with message passing not making use of the relation type.

By only using the connectivity information, but not the relation type information, this module can utilize message passing layers defined on uni-relational graphs, which are the majority of available layers from the PyTorch Geometric library.

Here, we create a two-layer torch_geometric.nn.conv.GCNConv on top of an pykeen.nn.representation.Embedding:

from pykeen.datasets import get_dataset

embedding_dim = 64
dataset = get_dataset(dataset="nations")
r = SimpleMessagePassingRepresentation(
    triples_factory=dataset.training,
    base_kwargs=dict(shape=embedding_dim),
    layers=["gcn"] * 2,
    layers_kwargs=dict(in_channels=embedding_dim, out_channels=embedding_dim),
)

Initialize the representation.

Parameters:

  • triples_factory (CoreTriplesFactory) – the factory comprising the training triples used for message passing

  • layers (Sequence[None]) – the message passing layer(s) or hints thereof

  • layers_kwargs (Mapping[str, Any] | None | Sequence[Mapping[str, Any] | None]) – additional keyword-based parameters passed to the layers upon instantiation

  • base (str | Representation | type[Representation] | None) – the base representations for entities, or a hint thereof

  • base_kwargs (Mapping[str, Any] | None) – additional keyword-based parameters passed to the base representations upon instantiation

  • shape (tuple[int, ...]) – the output shape. Defaults to the base representation shape. Has to match to output shape of the last message passing layer.

  • max_id (int) – the number of representations. If provided, has to match the base representation’s max_id

  • activations (str | Module | type[Module] | None | Sequence[str | Module | type[Module] | None]) – the activation(s), or hints thereof

  • activations_kwargs (Mapping[str, Any] | None | Sequence[Mapping[str, Any] | None]) – additional keyword-based parameters passed to the activations upon instantiation

  • restrict_k_hop (bool) – whether to restrict the message passing only to the k-hop neighborhood, when only some indices are requested. This utilizes torch_geometric.utils.k_hop_subgraph().

  • kwargs – additional keyword-based parameters passed to Representation.__init__()

Raises:

  • ImportError – if PyTorch Geometric is not installed

  • ValueError – if the number of activations and message passing layers do not match (after input normalization)

Methods Summary

pass_messages(x, edge_index[, edge_mask])

Perform the message passing steps.

Methods Documentation

pass_messages(x: Tensor, edge_index: Tensor, edge_mask: Tensor | None = None) → Tensor[source]

Perform the message passing steps.

Parameters:

  • x (Tensor) – shape: (n, d_in) the base entity representations

  • edge_index (Tensor) – shape: (num_selected_edges,) the edge index (which may already be a selection of the full edge index)

  • edge_mask (Tensor | None) – shape: (num_edges,) an edge mask if message passing is restricted

Returns:

shape: (n, d_out) the enriched entity representations

Return type:

Tensor