k3-node 1.0.0__py3-none-any.whl

k3_node/models/pmlp.py

@@ -0,0 +1,157 @@
+from typing import Optional
+
+import keras
+from keras import ops
+
+from k3_node.layers.conv import SimpleConv
+from k3_node.layers.norm import BatchNorm
+
+
+class PMLP(keras.Model):
+    r"""The P(ropagational)MLP model from the `"Graph Neural Networks are
+    Inherently Good Generalizers: Insights by Bridging GNNs and MLPs"
+    <https://arxiv.org/abs/2212.09034>`_ paper.
+
+    :class:`PMLP` is identical to a standard MLP during training, but then
+    adopts a GNN architecture during testing.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        hidden_channels (int): Size of each hidden sample.
+        out_channels (int): Size of each output sample.
+        num_layers (int): The number of layers.
+        dropout (float, optional): Dropout probability of each hidden
+            embedding. (default: :obj:`0.`)
+        norm (bool, optional): If set to :obj:`False`, will not apply batch
+            normalization. (default: :obj:`True`)
+        bias (bool, optional): If set to :obj:`False`, the module will not
+            learn additive biases. (default: :obj:`True`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import PMLP
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+
+        model = PMLP(in_channels=8, hidden_channels=32, out_channels=4, num_layers=2)
+        print(tuple(model(x, edge_index).shape))  # (10, 4): message passing is used at inference
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        num_layers: int,
+        dropout: float = 0.,
+        norm: bool = True,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+        self.dropout_rate = dropout
+        self.use_norm = norm
+        self.use_bias = bias
+        # Instance-level training flag matching PyG's API: pmlp.training = False
+        self.training = True
+
+        # Build weight dimensions
+        dims = [in_channels] + [hidden_channels] * (num_layers - 1) + [out_channels]
+        self._weight_shapes = [(dims[i], dims[i + 1]) for i in range(num_layers)]
+
+        # We store weights as raw keras Variables so they are always available
+        self._weights_list = []
+        self._biases_list = []
+        for i, (in_d, out_d) in enumerate(self._weight_shapes):
+            w = self.add_weight(
+                shape=(in_d, out_d),
+                initializer=keras.initializers.GlorotUniform(),
+                trainable=True,
+                name=f"weight_{i}",
+            )
+            self._weights_list.append(w)
+            if bias:
+                b = self.add_weight(
+                    shape=(out_d,),
+                    initializer="zeros",
+                    trainable=True,
+                    name=f"bias_{i}",
+                )
+                self._biases_list.append(b)
+            else:
+                self._biases_list.append(None)
+
+        self._norm_layer = None
+        if norm:
+            self._norm_layer = BatchNorm(
+                hidden_channels,
+                affine=False,
+                track_running_stats=False,
+            )
+
+        self.conv = SimpleConv(aggr='mean', combine_root='self_loop')
+        self._dropout = keras.layers.Dropout(dropout)
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def reset_parameters(self) -> None:
+        r"""Resets all learnable parameters of the module."""
+        for i, (in_d, out_d) in enumerate(self._weight_shapes):
+            self._weights_list[i].assign(
+                keras.initializers.GlorotUniform()(shape=(in_d, out_d))
+            )
+            if self.use_bias and self._biases_list[i] is not None:
+                self._biases_list[i].assign(ops.zeros((out_d,)))
+
+    def call(self, x, edge_index=None, training=None):
+        """Forward pass.
+
+        Args:
+            x (Tensor): The node features of shape ``[N, in_channels]``.
+            edge_index (Tensor, optional): The edge indices. Required during
+                inference. (default: :obj:`None`)
+            training (bool, optional): Override the instance-level
+                ``self.training`` flag. (default: :obj:`None`)
+        """
+        if edge_index is None and isinstance(x, (tuple, list)):
+            if len(x) >= 2:
+                x, edge_index = x[0], x[1]
+        # Respect both call-time kwarg and instance-level flag (PyG compat)
+        is_training = training if training is not None else self.training
+
+        if not is_training and edge_index is None:
+            raise ValueError(
+                f"'edge_index' needs to be present during inference "
+                f"in '{self.__class__.__name__}'"
+            )
+
+        for i in range(self.num_layers):
+            # Apply weight multiplication (like x @ W^T where W is [in, out])
+            x = x @ self._weights_list[i]
+
+            if not is_training:
+                x = self.conv(x, edge_index)
+
+            if self.use_bias and self._biases_list[i] is not None:
+                x = x + self._biases_list[i]
+
+            if i != self.num_layers - 1:
+                if self._norm_layer is not None:
+                    x = self._norm_layer(x, training=is_training)
+                x = ops.relu(x)
+                x = self._dropout(x, training=is_training)
+
+        return x
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, num_layers={self.num_layers})')

k3_node/models/polynormer.py

@@ -0,0 +1,229 @@
+from typing import Optional
+
+import keras
+from keras import ops
+
+from k3_node.layers.aggr.base import from_dense_batch, to_dense_batch
+import numpy as np
+
+from k3_node.layers.conv import GATConv, GCNConv
+from k3_node.layers.attention import PolynormerAttention
+
+
+class Polynormer(keras.layers.Layer):
+    r"""The Polynormer module from the `"Polynormer: polynomial-expressive
+    graph transformer in linear time"
+    <https://arxiv.org/abs/2403.01232>`_ paper.
+
+    Args:
+        in_channels (int): Input channels.
+        hidden_channels (int): Hidden channels.
+        out_channels (int): Output channels.
+        local_layers (int): The number of local attention layers.
+            (default: :obj:`7`)
+        global_layers (int): The number of global attention layers.
+            (default: :obj:`2`)
+        in_dropout (float): Input dropout rate.
+            (default: :obj:`0.15`)
+        dropout (float): Dropout rate.
+            (default: :obj:`0.5`)
+        global_dropout (float): Global dropout rate.
+            (default: :obj:`0.5`)
+        heads (int): The number of heads.
+            (default: :obj:`1`)
+        beta (float): Aggregate type.
+            (default: :obj:`0.9`)
+        qk_shared (bool, optional): Whether weight of query and key are shared.
+            (default: :obj:`True`)
+        pre_ln (bool): Pre layer normalization.
+            (default: :obj:`False`)
+        post_bn (bool): Post batch normalization.
+            (default: :obj:`True`)
+        local_attn (bool): Whether use local attention (GATConv vs GCNConv).
+            (default: :obj:`False`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import Polynormer
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+
+        batch = np.repeat([0, 1], 5)  # two graphs with 5 nodes each
+        model = Polynormer(in_channels=8, hidden_channels=32, out_channels=4,
+                           local_layers=2, global_layers=1, heads=2)
+        out = model(x, edge_index, batch)
+        print(tuple(out.shape))  # (10, 4)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        local_layers: int = 7,
+        global_layers: int = 2,
+        in_dropout: float = 0.15,
+        dropout: float = 0.5,
+        global_dropout: float = 0.5,
+        heads: int = 1,
+        beta: float = 0.9,
+        qk_shared: bool = False,
+        pre_ln: bool = False,
+        post_bn: bool = True,
+        local_attn: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+
+        self._global = False
+        self.in_drop = in_dropout
+        self.dropout = dropout
+        self.pre_ln = pre_ln
+        self.post_bn = post_bn
+        self.beta = beta
+        self.heads = heads
+        self.hidden_channels = hidden_channels
+        self.local_attn = local_attn
+
+        inner_channels = heads * hidden_channels
+
+        self.h_lins = []
+        self.local_convs = []
+        self.lins = []
+        self.lns = []
+        self.pre_lns = [] if pre_ln else None
+        self.post_bns = [] if post_bn else None
+
+        # ---- First local layer ----
+        self.h_lins.append(keras.layers.Dense(inner_channels))
+        if local_attn:
+            self.local_convs.append(
+                GATConv(in_channels, hidden_channels, heads=heads, concat=True,
+                        add_self_loops=False, bias=False)
+            )
+        else:
+            self.local_convs.append(
+                GCNConv(in_channels, inner_channels, cached=False, normalize=True)
+            )
+        self.lins.append(keras.layers.Dense(inner_channels))
+        self.lns.append(keras.layers.LayerNormalization(epsilon=1e-5))
+        if pre_ln:
+            self.pre_lns.append(keras.layers.LayerNormalization(epsilon=1e-5))
+        if post_bn:
+            self.post_bns.append(
+                keras.layers.BatchNormalization(
+                    center=True, scale=True, momentum=0.9, epsilon=1e-5,
+                )
+            )
+
+        # ---- Subsequent local layers ----
+        for _ in range(local_layers - 1):
+            self.h_lins.append(keras.layers.Dense(inner_channels))
+            if local_attn:
+                self.local_convs.append(
+                    GATConv(inner_channels, hidden_channels, heads=heads,
+                            concat=True, add_self_loops=False, bias=False)
+                )
+            else:
+                self.local_convs.append(
+                    GCNConv(inner_channels, inner_channels, cached=False,
+                            normalize=True)
+                )
+            self.lins.append(keras.layers.Dense(inner_channels))
+            self.lns.append(keras.layers.LayerNormalization(epsilon=1e-5))
+            if pre_ln:
+                self.pre_lns.append(keras.layers.LayerNormalization(epsilon=1e-5))
+            if post_bn:
+                self.post_bns.append(
+                    keras.layers.BatchNormalization(
+                        center=True, scale=True, momentum=0.9, epsilon=1e-5,
+                    )
+                )
+
+        self.lin_in = keras.layers.Dense(inner_channels)
+        self.ln = keras.layers.LayerNormalization(epsilon=1e-5)
+
+        self.global_attn = [
+            PolynormerAttention(
+                channels=hidden_channels,
+                heads=heads,
+                head_channels=hidden_channels,
+                beta=beta,
+                dropout=global_dropout,
+                qk_shared=qk_shared,
+            )
+            for _ in range(global_layers)
+        ]
+
+        self.pred_local = keras.layers.Dense(out_channels)
+        self.pred_global = keras.layers.Dense(out_channels)
+
+        self._in_dropout = keras.layers.Dropout(in_dropout)
+        self._dropout = keras.layers.Dropout(dropout)
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def reset_parameters(self) -> None:
+        r"""Resets all learnable parameters of the module."""
+        # Keras layers reinitialize on next forward; no-op for unbuilt layers.
+        pass
+
+    def call(self, x, edge_index, batch: Optional[object] = None, training=None):
+        r"""Forward pass.
+
+        Args:
+            x (Tensor): The input node features.
+            edge_index (Tensor): The edge indices.
+            batch (Tensor, optional): The batch vector assigning each node to
+                a graph. (default: :obj:`None`)
+            training (bool, optional): Whether in training mode.
+                (default: :obj:`None`)
+        """
+        x = self._in_dropout(x, training=training)
+
+        # ---- Equivariant local attention ----
+        x_local = 0
+        for i, local_conv in enumerate(self.local_convs):
+            if self.pre_ln:
+                x = self.pre_lns[i](x)
+            h = self.h_lins[i](x)
+            h = ops.relu(h)
+            x = local_conv(x, edge_index) + self.lins[i](x)
+            if self.post_bn:
+                x = self.post_bns[i](x, training=training)
+            x = ops.relu(x)
+            x = self._dropout(x, training=training)
+            x = (1 - self.beta) * self.lns[i](h * x) + self.beta * x
+            x_local = x_local + x
+
+        # ---- Equivariant global attention ----
+        if self._global:
+            # Sort nodes by batch assignment (required by to_dense_batch)
+            batch_i = ops.cast(batch, "int32")
+            indices = ops.argsort(batch_i)
+            rev_perm = ops.argsort(indices)
+            batch_sorted = ops.take(batch_i, indices, axis=0)
+            x_local_sorted = self.ln(ops.take(x_local, indices, axis=0))
+
+            x_global, mask = to_dense_batch(x_local_sorted, batch_sorted)
+            for attn in self.global_attn:
+                x_global = attn(x_global, mask=mask, training=training)
+
+            # Flatten and undo the sort
+            x = ops.take(from_dense_batch(x_global, batch_sorted), rev_perm, axis=0)
+            x = self.pred_global(x)
+        else:
+            x = self.pred_local(x_local)
+
+        return ops.log_softmax(x, axis=-1)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}('
+                f'in_channels={self.hidden_channels}, '
+                f'hidden_channels={self.hidden_channels}, '
+                f'heads={self.heads})')
+

k3_node/models/rect.py

@@ -0,0 +1,93 @@
+from typing import Optional
+
+import keras
+from keras import ops
+
+from k3_node.layers.conv import GCNConv
+from k3_node.layers.conv.utils import scatter
+
+
+class RECT_L(keras.layers.Layer):
+    r"""The RECT model, *i.e.* its supervised RECT-L part, from the
+    `"Network Embedding with Completely-imbalanced Labels"
+    <https://arxiv.org/abs/2007.03545>`_ paper.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        hidden_channels (int): Intermediate size of each sample.
+        normalize (bool, optional): Whether to add self-loops and compute
+            symmetric normalization coefficients on-the-fly.
+            (default: :obj:`True`)
+        dropout (float, optional): The dropout probability.
+            (default: :obj:`0.0`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import RECT_L
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+
+        model = RECT_L(in_channels=8, hidden_channels=16)
+        out = model(x, edge_index)  # reconstructs the (semantic) input features
+        print(tuple(out.shape))  # (10, 8)
+        print(tuple(model.embed(x, edge_index).shape))  # (10, 16): node embeddings
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        normalize: bool = True,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.dropout = dropout
+
+        self.conv = GCNConv(in_channels, hidden_channels, normalize=normalize)
+        self.lin = keras.layers.Dense(in_channels)
+        self._dropout = keras.layers.Dropout(dropout) if dropout > 0 else None
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        self.conv.reset_parameters()
+        if self.lin.built:
+            self.lin.kernel.assign(
+                keras.initializers.GlorotUniform()(self.lin.kernel.shape)
+            )
+            if self.lin.bias is not None:
+                self.lin.bias.assign(ops.zeros(self.lin.bias.shape))
+
+    def call(self, x, edge_index, edge_weight=None, training=None):
+        x = self.conv(x, edge_index, edge_weight=edge_weight)
+        if self._dropout is not None:
+            x = self._dropout(x, training=training)
+        return self.lin(x)
+
+    def embed(self, x, edge_index, edge_weight=None):
+        return self.conv(x, edge_index, edge_weight=edge_weight)
+
+    def get_semantic_labels(self, x, y, mask):
+        r"""Replaces the original labels by their class-centers."""
+        mask_shape = ops.shape(mask)
+        if len(mask_shape) == 1 and 'bool' in str(mask.dtype):
+            y_sub = y[mask]
+            x_sub = x[mask]
+        else:
+            y_sub = ops.take(y, mask, axis=0)
+            x_sub = ops.take(x, mask, axis=0)
+
+        num_classes = int(ops.max(y_sub)) + 1
+        mean = scatter(x_sub, y_sub, dim=0, dim_size=num_classes, reduce='mean')
+        return ops.take(mean, y_sub, axis=0)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.hidden_channels})')

k3_node/models/renet.py

@@ -0,0 +1,221 @@
+from typing import Callable, List, Optional, Tuple
+import math
+import numpy as np
+import keras
+from keras import ops
+
+from k3_node.layers.aggr import MeanAggregation
+from k3_node.ops.creation import repeat
+
+
+class RENet(keras.Model):
+    r"""The Recurrent Event Network model from the `"Recurrent Event Network
+    for Reasoning over Temporal Knowledge Graphs"
+    <https://arxiv.org/abs/1904.05530>`_ paper.
+
+    Args:
+        num_nodes (int): The number of nodes in the knowledge graph.
+        num_rels (int): The number of relations in the knowledge graph.
+        hidden_channels (int): Hidden size of node and relation embeddings.
+        seq_len (int): The sequence length of past events.
+        num_layers (int, optional): The number of recurrent layers.
+            (default: :obj:`1`)
+        dropout (float, optional): Dropout rate before final prediction.
+            (default: :obj:`0.0`)
+        bias (bool, optional): If set to :obj:`False`, all layers will not
+            learn an additive bias. (default: :obj:`True`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import RENet
+
+        model = RENet(num_nodes=5, num_rels=4, hidden_channels=16, seq_len=3)
+        sub, rel, obj = np.array([0, 1]), np.array([0, 1]), np.array([2, 3])  # queries at time t
+        # Neighbor histories of subjects and objects: neighbor id, timestep and query index
+        h_sub, h_sub_t, h_sub_batch = np.array([0, 1, 2]), np.array([0, 1, 0]), np.array([0, 0, 1])
+        h_obj, h_obj_t, h_obj_batch = np.array([1, 2, 3]), np.array([1, 2, 0]), np.array([0, 0, 1])
+        log_prob_obj, log_prob_sub = model(sub, rel, obj, h_sub, h_sub_t, h_sub_batch, h_obj, h_obj_t, h_obj_batch)
+        print(tuple(log_prob_obj.shape))  # (2, 5): scores over all entities for each query
+        ```
+    """
+    def __init__(
+        self,
+        num_nodes: int,
+        num_rels: int,
+        hidden_channels: int,
+        seq_len: int,
+        num_layers: int = 1,
+        dropout: float = 0.0,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.num_nodes = num_nodes
+        self.num_rels = num_rels
+        self.hidden_channels = hidden_channels
+        self.seq_len = seq_len
+        self.dropout_rate = dropout
+        self.num_layers = num_layers
+
+        self.ent = self.add_weight(
+            name="ent",
+            shape=(num_nodes, hidden_channels),
+            initializer=keras.initializers.GlorotUniform(),
+        )
+        self.rel = self.add_weight(
+            name="rel",
+            shape=(num_rels, hidden_channels),
+            initializer=keras.initializers.GlorotUniform(),
+        )
+
+        self.sub_gru = keras.layers.GRU(
+            hidden_channels,
+            return_sequences=False,
+            use_bias=bias,
+        )
+        self.obj_gru = keras.layers.GRU(
+            hidden_channels,
+            return_sequences=False,
+            use_bias=bias,
+        )
+
+        self.sub_lin = keras.layers.Dense(num_nodes, use_bias=bias)
+        self.obj_lin = keras.layers.Dense(num_nodes, use_bias=bias)
+        self.drop = keras.layers.Dropout(dropout)
+        self.mean_aggr = MeanAggregation()
+
+    def reset_parameters(self):
+        self.ent.assign(
+            keras.initializers.GlorotUniform()(self.ent.shape)
+        )
+        self.rel.assign(
+            keras.initializers.GlorotUniform()(self.rel.shape)
+        )
+
+    def build(self, input_shape=None):
+        self.built = True
+
+
+    def call(
+        self,
+        sub,
+        rel,
+        obj,
+        h_sub,
+        h_sub_t,
+        h_sub_batch,
+        h_obj,
+        h_obj_t,
+        h_obj_batch,
+        training=False,
+    ):
+        batch_size = ops.shape(sub)[0]
+        seq_len = self.seq_len
+
+        h_sub_t = h_sub_t + h_sub_batch * seq_len
+        h_obj_t = h_obj_t + h_obj_batch * seq_len
+
+        ent_h_sub = ops.take(self.ent, h_sub, axis=0)
+        ent_h_obj = ops.take(self.ent, h_obj, axis=0)
+
+        h_sub_scatter = self.mean_aggr(
+            ent_h_sub, index=h_sub_t, dim_size=batch_size * seq_len, dim=0
+        )
+        h_sub = ops.reshape(h_sub_scatter, (-1, seq_len, self.hidden_channels))  # static feature size
+
+        h_obj_scatter = self.mean_aggr(
+            ent_h_obj, index=h_obj_t, dim_size=batch_size * seq_len, dim=0
+        )
+        h_obj = ops.reshape(h_obj_scatter, (-1, seq_len, self.hidden_channels))  # static feature size
+
+        sub_emb = ops.take(self.ent, sub, axis=0)
+        rel_emb = ops.take(self.rel, rel, axis=0)
+        obj_emb = ops.take(self.ent, obj, axis=0)
+
+        sub_rep = repeat(ops.expand_dims(sub_emb, 1), seq_len, axis=1)
+        rel_rep = repeat(ops.expand_dims(rel_emb, 1), seq_len, axis=1)
+        obj_rep = repeat(ops.expand_dims(obj_emb, 1), seq_len, axis=1)
+
+        gru_sub_in = ops.concatenate([sub_rep, h_sub, rel_rep], axis=-1)
+        gru_obj_in = ops.concatenate([obj_rep, h_obj, rel_rep], axis=-1)
+
+        h_sub = self.sub_gru(gru_sub_in, training=training)
+        h_obj = self.obj_gru(gru_obj_in, training=training)
+
+        h_sub = ops.concatenate([sub_emb, h_sub, rel_emb], axis=-1)
+        h_obj = ops.concatenate([obj_emb, h_obj, rel_emb], axis=-1)
+
+        h_sub = self.drop(h_sub, training=training)
+        h_obj = self.drop(h_obj, training=training)
+
+        log_prob_obj = ops.log_softmax(self.sub_lin(h_sub), axis=-1)
+        log_prob_sub = ops.log_softmax(self.obj_lin(h_obj), axis=-1)
+
+        return log_prob_obj, log_prob_sub
+
+    @staticmethod
+    def pre_transform(seq_len: int) -> Callable:
+        r"""Returns a pre-transform that adds to every event (processed in time order) the history
+        of its subject and object: the entities they were linked to by the same relation in each of
+        the last ``seq_len`` time steps (``h_sub`` / ``h_obj``, with the step in ``h_sub_t`` /
+        ``h_obj_t``), as in PyG."""
+
+        class PreTransform:
+            def __init__(self, seq_len):
+                self.seq_len = seq_len
+                self.t_last = 0
+                self.sub_hist, self.obj_hist = {}, {}  # node -> list of seq_len + 1 steps of (node, rel)
+
+            def _hist(self, hist, node):
+                if node not in hist:
+                    hist[node] = [[] for _ in range(self.seq_len + 1)]
+                return hist[node]
+
+            def _history(self, hist, node, rel):
+                steps = self._hist(hist, node)
+                nodes, ts = [], []
+                for s in range(self.seq_len):
+                    for other, r in steps[s]:
+                        if r == rel:
+                            nodes.append(other)
+                            ts.append(s)
+                return np.array(nodes, dtype=np.int64), np.array(ts, dtype=np.int64)
+
+            def __call__(self, data):
+                sub, rel, obj, t = int(data.sub), int(data.rel), int(data.obj), int(data.t)
+                if t > self.t_last:  # a new time step: forget the oldest one
+                    for hist in (self.sub_hist, self.obj_hist):
+                        for steps in hist.values():
+                            steps.pop(0)
+                            steps.append([])
+                    self.t_last = t
+                data.h_sub, data.h_sub_t = self._history(self.sub_hist, sub, rel)
+                data.h_obj, data.h_obj_t = self._history(self.obj_hist, obj, rel)
+                self._hist(self.sub_hist, sub)[-1].append((obj, rel))
+                self._hist(self.obj_hist, obj)[-1].append((sub, rel))
+                return data
+
+            def __repr__(self):
+                return f"{self.__class__.__name__}(seq_len={self.seq_len})"
+
+        return PreTransform(seq_len)
+
+    def test(self, logits, y):
+        r"""Given ground-truth :obj:`y`, computes Mean Reciprocal Rank (MRR)
+        and Hits at 1/3/10.
+        """
+        logits_np = ops.convert_to_numpy(logits)
+        y_np = ops.convert_to_numpy(y).reshape(-1, 1)
+
+        perm = np.argsort(-logits_np, axis=1)
+        mask = y_np == perm
+
+        rows, cols = np.nonzero(mask)
+        mrr = float(np.mean(1.0 / (cols + 1.0)))
+        hits1 = float(np.sum(cols < 1) / len(y_np))
+        hits3 = float(np.sum(cols < 3) / len(y_np))
+        hits10 = float(np.sum(cols < 10) / len(y_np))
+
+        return ops.convert_to_tensor([mrr, hits1, hits3, hits10], dtype="float32")