FeaturizedMessagePassingRepresentation

class FeaturizedMessagePassingRepresentation(triples_factory: CoreTriplesFactory, relation_representation: str | Representation | type[Representation] | None = None, relation_representation_kwargs: Mapping[str, Any] | None = None, relation_transformation: Module | None = None, **kwargs)

Bases: TypedMessagePassingRepresentation

A representation with message passing with uses edge features obtained from relation representations.

It (re-)uses a representation layer for relations to obtain edge features, which are then utilized by appropriate message passing layers, e.g., torch_geometric.nn.conv.GMMConv, or torch_geometric.nn.conv.GATConv. We further allow a (shared) transformation of edge features between layers.

The following example creates a two-layer GAT on top of the base representations:

from pykeen.datasets import get_dataset

embedding_dim = 64
dataset = get_dataset(dataset="nations")
r = FeaturizedMessagePassingRepresentation(
    triples_factory=dataset.training,
    base_kwargs=dict(shape=embedding_dim),
    relation_representation_kwargs=dict(
        shape=embedding_dim,
    ),
    layers="gat",
    layers_kwargs=dict(
        in_channels=embedding_dim,
        out_channels=embedding_dim,
        edge_dim=embedding_dim,  # should match relation dim
    ),
)

Initialize the representation.

Parameters:

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

• relation_representation (Representation) – the base representations for relations, or a hint thereof

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

• relation_transformation (Module | None) – an optional transformation to apply to the relation representations after each message passing step. If None, do not modify the representations.

• kwargs – additional keyword-based parameters passed to TypedMessagePassingRepresentation.__init__(), except the triples_factory

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

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