k3-node 1.0.0__py3-none-any.whl

k3_node/layers/kge/base.py

@@ -0,0 +1,255 @@
+from typing import Tuple
+
+import numpy as np
+import keras
+from keras import ops
+
+from k3_node.layers.kge.loader import KGTripletLoader
+from k3_node.training import no_grad
+
+try:
+    from tqdm import tqdm
+except ImportError:  # pragma: no cover
+    def tqdm(iterable, *args, **kwargs):
+        return iterable
+
+
+def normalize(x, p: float = 2.0, axis: int = -1, eps: float = 1e-12):
+    r"""Lp-normalizes `x` along `axis`, mirroring `torch.nn.functional.normalize`."""
+    norm = ops.power(ops.sum(ops.power(ops.abs(x), p), axis=axis, keepdims=True), 1.0 / p)
+    return x / ops.maximum(norm, eps)
+
+
+def margin_ranking_loss(pos_score, neg_score, margin: float = 1.0):
+    r"""Mirrors `torch.nn.functional.margin_ranking_loss` with `target=1`."""
+    return ops.mean(ops.relu(margin - pos_score + neg_score))
+
+
+def binary_cross_entropy_with_logits(logits, target):
+    r"""Numerically stable sigmoid cross-entropy, mirroring
+    `torch.nn.functional.binary_cross_entropy_with_logits`."""
+    zeros = ops.zeros_like(logits)
+    cond = logits >= zeros
+    relu_logits = ops.where(cond, logits, zeros)
+    neg_abs_logits = ops.where(cond, -logits, logits)
+    loss = relu_logits - logits * target + ops.log(1.0 + ops.exp(neg_abs_logits))
+    return ops.mean(loss)
+
+
+class KGEModel(keras.layers.Layer):
+    r"""An abstract base class for implementing custom KGE models.
+
+    Args:
+        num_nodes (int): The number of nodes/entities in the graph.
+        num_relations (int): The number of relations in the graph.
+        hidden_channels (int): The hidden embedding size.
+        sparse (bool, optional): Kept for API compatibility with PyG; has no
+            effect since Keras optimizers do not distinguish sparse
+            embedding gradients the way PyTorch does. (default: :obj:`False`)
+    """
+    def __init__(
+        self,
+        num_nodes: int,
+        num_relations: int,
+        hidden_channels: int,
+        sparse: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.num_nodes = num_nodes
+        self.num_relations = num_relations
+        self.hidden_channels = hidden_channels
+        self.sparse = sparse
+
+        self.node_emb = keras.layers.Embedding(num_nodes, hidden_channels)
+        self.rel_emb = keras.layers.Embedding(num_relations, hidden_channels)
+        self.node_emb.build((None,))
+        self.rel_emb.build((None,))
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        self.node_emb.embeddings.assign(
+            self.node_emb.embeddings_initializer(ops.shape(self.node_emb.embeddings))
+        )
+        self.rel_emb.embeddings.assign(
+            self.rel_emb.embeddings_initializer(ops.shape(self.rel_emb.embeddings))
+        )
+
+    def call(self, head_index, rel_type, tail_index):
+        r"""Returns the score for the given triplet.
+
+        Args:
+            head_index: The head indices.
+            rel_type: The relation type.
+            tail_index: The tail indices.
+        """
+        raise NotImplementedError
+
+    def loss(self, head_index, rel_type, tail_index):
+        r"""Returns the loss value for the given triplet."""
+        raise NotImplementedError
+
+    def loader(self, head_index, rel_type, tail_index, **kwargs):
+        r"""Returns a mini-batch loader that samples a subset of triplets.
+
+        Args:
+            head_index: The head indices.
+            rel_type: The relation type.
+            tail_index: The tail indices.
+            **kwargs (optional): Additional arguments of
+                :class:`k3_node.layers.kge.KGTripletLoader`, such as
+                `batch_size`, `shuffle` or `drop_last`.
+        """
+        return KGTripletLoader(head_index, rel_type, tail_index, **kwargs)
+
+    def test(
+        self,
+        head_index,
+        rel_type,
+        tail_index,
+        batch_size: int,
+        k: int = 10,
+        log: bool = True,
+    ) -> Tuple[float, float, float]:
+        r"""Evaluates the model quality by computing Mean Rank, MRR and
+        Hits@:math:`k` across all possible tail entities.
+
+        Args:
+            head_index: The head indices.
+            rel_type: The relation type.
+            tail_index: The tail indices.
+            batch_size (int): The batch size to use for evaluating.
+            k (int, optional): The :math:`k` in Hits @ :math:`k`.
+                (default: :obj:`10`)
+            log (bool, optional): If set to :obj:`False`, will not print a
+                progress bar to the console. (default: :obj:`True`)
+        """
+        head_index = ops.convert_to_numpy(head_index)
+        rel_type = ops.convert_to_numpy(rel_type)
+        tail_index = ops.convert_to_numpy(tail_index)
+
+        arange = range(head_index.shape[0])
+        arange = tqdm(arange) if log else arange
+
+        mean_ranks, reciprocal_ranks, hits_at_k = [], [], []
+        for i in arange:
+            h, r, t = int(head_index[i]), int(rel_type[i]), int(tail_index[i])
+
+            scores = []
+            tail_indices = np.arange(self.num_nodes)
+            for start in range(0, self.num_nodes, batch_size):
+                ts = tail_indices[start:start + batch_size]
+                hs = np.full_like(ts, h)
+                rs = np.full_like(ts, r)
+                with no_grad():  # as PyG's @torch.no_grad()
+                    out = self(
+                        ops.convert_to_tensor(hs),
+                        ops.convert_to_tensor(rs),
+                        ops.convert_to_tensor(ts),
+                    )
+                scores.append(ops.convert_to_numpy(out))
+            scores = np.concatenate(scores)
+            rank = int(np.nonzero(np.argsort(-scores) == t)[0][0])
+
+            mean_ranks.append(rank)
+            reciprocal_ranks.append(1.0 / (rank + 1))
+            hits_at_k.append(rank < k)
+
+        mean_rank = float(np.mean(mean_ranks))
+        mrr = float(np.mean(reciprocal_ranks))
+        hits_at_k = float(np.mean(hits_at_k))
+
+        return mean_rank, mrr, hits_at_k
+
+    # ---- Keras-style training -------------------------------------------------------------------
+    def compile(self, optimizer):
+        r"""Sets the optimizer used by :meth:`fit`."""
+        self.optimizer = optimizer
+
+    def fit(self, data, epochs: int = 1, batch_size: int = 1000, validation_data=None,
+            validation_batch_size: int = 20000, verbose: int = 1):
+        r"""Trains on the triplets of ``data`` (``edge_index`` holds heads and tails, ``edge_type``
+        the relations), in shuffled mini-batches.
+
+        Args:
+            data (Data): The training knowledge graph.
+            epochs (int): Passes over all triplets. (default: ``1``)
+            batch_size (int): Triplets per gradient step. (default: ``1000``)
+            validation_data (Data, optional): Evaluated with :meth:`evaluate` after every epoch
+                (this ranks all entities for every triplet, so keep it small).
+            validation_batch_size (int): Entities scored at once during validation.
+            verbose (int): ``0`` is silent, otherwise one line is printed per epoch.
+
+        Returns:
+            dict: The mean loss (and validation metrics) of every epoch.
+        """
+        from k3_node.training import gradient_step
+
+        if getattr(self, "optimizer", None) is None:
+            raise ValueError("Call `compile(optimizer=...)` before `fit`.")
+        head, tail = data.edge_index[0], data.edge_index[1]
+        loader = self.loader(head, data.edge_type, tail, batch_size=batch_size, shuffle=True)
+        history = {"loss": []}
+        for epoch in range(1, epochs + 1):
+            total = count = 0
+            for h, r, t in loader:
+                loss = gradient_step(lambda: self.loss(h, r, t), self.trainable_variables, self.optimizer)
+                total += loss * int(h.shape[0])
+                count += int(h.shape[0])
+            logs = {"loss": total / count}
+            if validation_data is not None:
+                logs.update({f"val_{k}": v for k, v in
+                             self.evaluate(validation_data, batch_size=validation_batch_size).items()})
+            for key, value in logs.items():
+                history.setdefault(key, []).append(value)
+            if verbose:
+                print(f"Epoch {epoch:03d}: " + ", ".join(f"{k}: {v:.4f}" for k, v in logs.items()))
+        return history
+
+    def evaluate(self, data, batch_size: int = 20000, k: int = 10):
+        r"""Ranks the true tail of every triplet in ``data`` among all entities and returns the
+        mean rank, the mean reciprocal rank (MRR) and Hits@``k``."""
+        mean_rank, mrr, hits = self.test(data.edge_index[0], data.edge_type, data.edge_index[1],
+                                         batch_size=batch_size, k=k, log=False)
+        return {"mean_rank": mean_rank, "mrr": mrr, f"hits@{k}": hits}
+
+    def random_sample(
+        self,
+        head_index,
+        rel_type,
+        tail_index,
+    ):
+        r"""Randomly samples negative triplets by either replacing the head or
+        the tail (but not both).
+
+        Args:
+            head_index: The head indices.
+            rel_type: The relation type.
+            tail_index: The tail indices.
+        """
+        num_triplets = ops.shape(head_index)[0]
+        num_negatives = num_triplets // 2
+
+        rnd_index = keras.random.randint(
+            ops.shape(head_index), 0, self.num_nodes, dtype="int32"
+        )
+        rnd_index = ops.cast(rnd_index, head_index.dtype)
+
+        head_index = ops.concatenate(
+            [rnd_index[:num_negatives], head_index[num_negatives:]], axis=0
+        )
+        tail_index = ops.concatenate(
+            [tail_index[:num_negatives], rnd_index[num_negatives:]], axis=0
+        )
+
+        return head_index, rel_type, tail_index
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}({self.num_nodes}, "
+            f"num_relations={self.num_relations}, "
+            f"hidden_channels={self.hidden_channels})"
+        )
+
+    __str__ = __repr__

k3_node/layers/kge/complex.py

@@ -0,0 +1,98 @@
+import keras
+from keras import ops
+
+from k3_node.layers.kge.base import KGEModel, binary_cross_entropy_with_logits
+
+
+def triple_dot(x, y, z):
+    return ops.sum(x * y * z, axis=-1)
+
+
+class ComplEx(KGEModel):
+    r"""The ComplEx model from the `"Complex Embeddings for Simple Link
+    Prediction" <https://arxiv.org/abs/1606.06357>`_ paper.
+
+    :class:`ComplEx` models relations as complex-valued bilinear mappings
+    between head and tail entities using the Hermetian dot product.
+    The entities and relations are embedded in different dimensional spaces,
+    resulting in the scoring function:
+
+    .. math::
+        d(h, r, t) = Re(< \mathbf{e}_h,  \mathbf{e}_r, \mathbf{e}_t>)
+
+    Args:
+        num_nodes (int): The number of nodes/entities in the graph.
+        num_relations (int): The number of relations in the graph.
+        hidden_channels (int): The hidden embedding size.
+        sparse (bool, optional): Kept for API compatibility. (default: :obj:`False`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import ComplEx
+
+        head = np.random.randint(0, 20, size=(10,))  # 10 (head, relation, tail) triples
+        rel = np.random.randint(0, 5, size=(10,))
+        tail = np.random.randint(0, 20, size=(10,))
+
+        model = ComplEx(num_nodes=20, num_relations=5, hidden_channels=8)
+        score = model(head, rel, tail)  # plausibility score of every triple
+        print(tuple(score.shape))  # (10,)
+        loss = model.loss(head, rel, tail)  # training loss against randomly corrupted triples
+        print(tuple(loss.shape))  # (): a scalar
+        ```
+    """
+    def __init__(
+        self,
+        num_nodes: int,
+        num_relations: int,
+        hidden_channels: int,
+        sparse: bool = False,
+        **kwargs,
+    ):
+        super().__init__(num_nodes, num_relations, hidden_channels, sparse, **kwargs)
+
+        self.node_emb_im = keras.layers.Embedding(num_nodes, hidden_channels)
+        self.rel_emb_im = keras.layers.Embedding(num_relations, hidden_channels)
+        self.node_emb_im.build((None,))
+        self.rel_emb_im.build((None,))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        # A new initializer per tensor: a reused unseeded Keras 3 initializer returns the same values on every call.
+        glorot = lambda shape: keras.initializers.GlorotUniform()(shape)
+        self.node_emb.embeddings.assign(glorot(ops.shape(self.node_emb.embeddings)))
+        self.node_emb_im.embeddings.assign(glorot(ops.shape(self.node_emb_im.embeddings)))
+        self.rel_emb.embeddings.assign(glorot(ops.shape(self.rel_emb.embeddings)))
+        self.rel_emb_im.embeddings.assign(glorot(ops.shape(self.rel_emb_im.embeddings)))
+
+    def call(self, head_index, rel_type, tail_index):
+        head_index = ops.cast(head_index, "int32")
+        rel_type = ops.cast(rel_type, "int32")
+        tail_index = ops.cast(tail_index, "int32")
+
+        head_re = self.node_emb(head_index)
+        head_im = self.node_emb_im(head_index)
+        rel_re = self.rel_emb(rel_type)
+        rel_im = self.rel_emb_im(rel_type)
+        tail_re = self.node_emb(tail_index)
+        tail_im = self.node_emb_im(tail_index)
+
+        return (
+            triple_dot(head_re, rel_re, tail_re)
+            + triple_dot(head_im, rel_re, tail_im)
+            + triple_dot(head_re, rel_im, tail_im)
+            - triple_dot(head_im, rel_im, tail_re)
+        )
+
+    def loss(self, head_index, rel_type, tail_index):
+        pos_score = self(head_index, rel_type, tail_index)
+        neg_score = self(*self.random_sample(head_index, rel_type, tail_index))
+        scores = ops.concatenate([pos_score, neg_score], axis=0)
+
+        pos_target = ops.ones_like(pos_score)
+        neg_target = ops.zeros_like(neg_score)
+        target = ops.concatenate([pos_target, neg_target], axis=0)
+
+        return binary_cross_entropy_with_logits(scores, target)

k3_node/layers/kge/distmult.py

@@ -0,0 +1,79 @@
+import keras
+from keras import ops
+
+from k3_node.layers.kge.base import KGEModel, margin_ranking_loss
+
+
+class DistMult(KGEModel):
+    r"""The DistMult model from the `"Embedding Entities and Relations for
+    Learning and Inference in Knowledge Bases"
+    <https://arxiv.org/abs/1412.6575>`_ paper.
+
+    :class:`DistMult` models relations as diagonal matrices, which simplifies
+    the bi-linear interaction between the head and tail entities to the score
+    function:
+
+    .. math::
+        d(h, r, t) = < \mathbf{e}_h,  \mathbf{e}_r, \mathbf{e}_t >
+
+    Args:
+        num_nodes (int): The number of nodes/entities in the graph.
+        num_relations (int): The number of relations in the graph.
+        hidden_channels (int): The hidden embedding size.
+        margin (float, optional): The margin of the ranking loss.
+            (default: :obj:`1.0`)
+        sparse (bool, optional): Kept for API compatibility. (default: :obj:`False`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import DistMult
+
+        head = np.random.randint(0, 20, size=(10,))  # 10 (head, relation, tail) triples
+        rel = np.random.randint(0, 5, size=(10,))
+        tail = np.random.randint(0, 20, size=(10,))
+
+        model = DistMult(num_nodes=20, num_relations=5, hidden_channels=8)
+        score = model(head, rel, tail)  # plausibility score of every triple
+        print(tuple(score.shape))  # (10,)
+        loss = model.loss(head, rel, tail)  # training loss against randomly corrupted triples
+        print(tuple(loss.shape))  # (): a scalar
+        ```
+    """
+    def __init__(
+        self,
+        num_nodes: int,
+        num_relations: int,
+        hidden_channels: int,
+        margin: float = 1.0,
+        sparse: bool = False,
+        **kwargs,
+    ):
+        super().__init__(num_nodes, num_relations, hidden_channels, sparse, **kwargs)
+
+        self.margin = margin
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        # A new initializer per tensor: a reused unseeded Keras 3 initializer returns the same values on every call.
+        glorot = lambda shape: keras.initializers.GlorotUniform()(shape)
+        self.node_emb.embeddings.assign(glorot(ops.shape(self.node_emb.embeddings)))
+        self.rel_emb.embeddings.assign(glorot(ops.shape(self.rel_emb.embeddings)))
+
+    def call(self, head_index, rel_type, tail_index):
+        head_index = ops.cast(head_index, "int32")
+        rel_type = ops.cast(rel_type, "int32")
+        tail_index = ops.cast(tail_index, "int32")
+
+        head = self.node_emb(head_index)
+        rel = self.rel_emb(rel_type)
+        tail = self.node_emb(tail_index)
+
+        return ops.sum(head * rel * tail, axis=-1)
+
+    def loss(self, head_index, rel_type, tail_index):
+        pos_score = self(head_index, rel_type, tail_index)
+        neg_score = self(*self.random_sample(head_index, rel_type, tail_index))
+
+        return margin_ranking_loss(pos_score, neg_score, margin=self.margin)

k3_node/layers/kge/loader.py

@@ -0,0 +1,50 @@
+import math
+
+import numpy as np
+from keras import ops
+
+
+class KGTripletLoader:
+    r"""A minimal, framework-agnostic batching iterator over knowledge-graph
+    triplets, mirroring :class:`torch.utils.data.DataLoader` usage in
+    :meth:`k3_node.layers.kge.KGEModel.loader`.
+
+    Args:
+        head_index: The head indices.
+        rel_type: The relation type.
+        tail_index: The tail indices.
+        batch_size (int, optional): The batch size. (default: :obj:`1`)
+        shuffle (bool, optional): If set to :obj:`True`, shuffles the
+            triplets at every epoch. (default: :obj:`False`)
+        drop_last (bool, optional): If set to :obj:`True`, drops the last
+            incomplete batch. (default: :obj:`False`)
+    """
+    def __init__(self, head_index, rel_type, tail_index, batch_size: int = 1,
+                shuffle: bool = False, drop_last: bool = False, **kwargs):
+        self.head_index = ops.convert_to_numpy(head_index)
+        self.rel_type = ops.convert_to_numpy(rel_type)
+        self.tail_index = ops.convert_to_numpy(tail_index)
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+        self.num_triplets = self.head_index.shape[0]
+
+    def __len__(self):
+        if self.drop_last:
+            return self.num_triplets // self.batch_size
+        return math.ceil(self.num_triplets / self.batch_size)
+
+    def __iter__(self):
+        indices = np.arange(self.num_triplets)
+        if self.shuffle:
+            np.random.shuffle(indices)
+
+        for start in range(0, self.num_triplets, self.batch_size):
+            batch_idx = indices[start:start + self.batch_size]
+            if self.drop_last and batch_idx.shape[0] < self.batch_size:
+                continue
+            yield (
+                ops.convert_to_tensor(self.head_index[batch_idx]),
+                ops.convert_to_tensor(self.rel_type[batch_idx]),
+                ops.convert_to_tensor(self.tail_index[batch_idx]),
+            )

k3_node/layers/kge/rotate.py

@@ -0,0 +1,103 @@
+import math
+
+import keras
+from keras import ops
+
+from k3_node.layers.kge.base import KGEModel, binary_cross_entropy_with_logits
+
+
+class RotatE(KGEModel):
+    r"""The RotatE model from the `"RotatE: Knowledge Graph Embedding by
+    Relational Rotation in Complex Space" <https://arxiv.org/abs/
+    1902.10197>`_ paper.
+
+    :class:`RotatE` models relations as a rotation in complex space
+    from head to tail such that
+
+    .. math::
+        \mathbf{e}_t = \mathbf{e}_h \circ \mathbf{e}_r,
+
+    resulting in the scoring function
+
+    .. math::
+        d(h, r, t) = - {\| \mathbf{e}_h \circ \mathbf{e}_r - \mathbf{e}_t \|}_p
+
+    Args:
+        num_nodes (int): The number of nodes/entities in the graph.
+        num_relations (int): The number of relations in the graph.
+        hidden_channels (int): The hidden embedding size.
+        margin (float, optional): The margin of the ranking loss.
+            (default: :obj:`1.0`)
+        sparse (bool, optional): Kept for API compatibility. (default: :obj:`False`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import RotatE
+
+        head = np.random.randint(0, 20, size=(10,))  # 10 (head, relation, tail) triples
+        rel = np.random.randint(0, 5, size=(10,))
+        tail = np.random.randint(0, 20, size=(10,))
+
+        model = RotatE(num_nodes=20, num_relations=5, hidden_channels=8)
+        score = model(head, rel, tail)  # plausibility score of every triple
+        print(tuple(score.shape))  # (10,)
+        loss = model.loss(head, rel, tail)  # training loss against randomly corrupted triples
+        print(tuple(loss.shape))  # (): a scalar
+        ```
+    """
+    def __init__(
+        self,
+        num_nodes: int,
+        num_relations: int,
+        hidden_channels: int,
+        margin: float = 1.0,
+        sparse: bool = False,
+        **kwargs,
+    ):
+        super().__init__(num_nodes, num_relations, hidden_channels, sparse, **kwargs)
+
+        self.margin = margin
+        self.node_emb_im = keras.layers.Embedding(num_nodes, hidden_channels)
+        self.node_emb_im.build((None,))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        # A new initializer per tensor: a reused unseeded Keras 3 initializer returns the same values on every call.
+        glorot = lambda shape: keras.initializers.GlorotUniform()(shape)
+        self.node_emb.embeddings.assign(glorot(ops.shape(self.node_emb.embeddings)))
+        self.node_emb_im.embeddings.assign(glorot(ops.shape(self.node_emb_im.embeddings)))
+        uniform = keras.initializers.RandomUniform(0, 2 * math.pi)
+        self.rel_emb.embeddings.assign(uniform(ops.shape(self.rel_emb.embeddings)))
+
+    def call(self, head_index, rel_type, tail_index):
+        head_index = ops.cast(head_index, "int32")
+        rel_type = ops.cast(rel_type, "int32")
+        tail_index = ops.cast(tail_index, "int32")
+
+        head_re = self.node_emb(head_index)
+        head_im = self.node_emb_im(head_index)
+        tail_re = self.node_emb(tail_index)
+        tail_im = self.node_emb_im(tail_index)
+
+        rel_theta = self.rel_emb(rel_type)
+        rel_re, rel_im = ops.cos(rel_theta), ops.sin(rel_theta)
+
+        re_score = (rel_re * head_re - rel_im * head_im) - tail_re
+        im_score = (rel_re * head_im + rel_im * head_re) - tail_im
+        complex_score = ops.stack([re_score, im_score], axis=2)
+        score = ops.sqrt(ops.sum(ops.square(complex_score), axis=(1, 2)))
+
+        return self.margin - score
+
+    def loss(self, head_index, rel_type, tail_index):
+        pos_score = self(head_index, rel_type, tail_index)
+        neg_score = self(*self.random_sample(head_index, rel_type, tail_index))
+        scores = ops.concatenate([pos_score, neg_score], axis=0)
+
+        pos_target = ops.ones_like(pos_score)
+        neg_target = ops.zeros_like(neg_score)
+        target = ops.concatenate([pos_target, neg_target], axis=0)
+
+        return binary_cross_entropy_with_logits(scores, target)

k3_node/layers/kge/test_kge.py

@@ -0,0 +1,76 @@
+from keras import ops
+
+from k3_node.layers.kge import TransE, DistMult, ComplEx, RotatE
+
+
+def _run_model(cls):
+    model = cls(num_nodes=10, num_relations=5, hidden_channels=32)
+    assert str(model) == f"{cls.__name__}(10, num_relations=5, hidden_channels=32)"
+
+    head_index = ops.convert_to_tensor([0, 2, 4, 6, 8])
+    rel_type = ops.convert_to_tensor([0, 1, 2, 3, 4])
+    tail_index = ops.convert_to_tensor([1, 3, 5, 7, 9])
+
+    loader = model.loader(head_index, rel_type, tail_index, batch_size=5)
+    for h, r, t in loader:
+        out = model(h, r, t)
+        assert ops.shape(out) == (5,)
+
+        loss = model.loss(h, r, t)
+        assert float(ops.convert_to_numpy(loss)) >= 0.0
+
+        mean_rank, mrr, hits = model.test(h, r, t, batch_size=5, log=False)
+        assert 0 <= mean_rank <= 10
+        assert 0 < mrr <= 1
+        assert hits == 1.0
+
+
+def test_transe():
+    _run_model(TransE)
+
+
+def test_distmult():
+    _run_model(DistMult)
+
+
+def test_complex():
+    _run_model(ComplEx)
+
+
+def test_rotate():
+    _run_model(RotatE)
+
+
+def test_complex_scoring():
+    model = ComplEx(num_nodes=5, num_relations=2, hidden_channels=1)
+
+    model.node_emb.embeddings.assign(
+        ops.convert_to_tensor([[2.0], [3.0], [5.0], [1.0], [2.0]], dtype="float32")
+    )
+    model.node_emb_im.embeddings.assign(
+        ops.convert_to_tensor([[4.0], [1.0], [3.0], [1.0], [2.0]], dtype="float32")
+    )
+    model.rel_emb.embeddings.assign(ops.convert_to_tensor([[2.0], [3.0]], dtype="float32"))
+    model.rel_emb_im.embeddings.assign(ops.convert_to_tensor([[3.0], [1.0]], dtype="float32"))
+
+    score = model(
+        ops.convert_to_tensor([1, 3]),
+        ops.convert_to_tensor([1, 0]),
+        ops.convert_to_tensor([2, 4]),
+    )
+    assert ops.convert_to_numpy(score).tolist() == [58.0, 8.0]
+
+
+def test_same_shape_embeddings_are_initialized_independently():
+    # A reused Keras 3 initializer instance repeats its values. For ComplEx, identical real and
+    # imaginary parts cancel the asymmetric term, collapsing the model to DistMult.
+    import numpy as np
+
+    complex_model = ComplEx(num_nodes=20, num_relations=5, hidden_channels=8)
+    rotate_model = RotatE(num_nodes=20, num_relations=5, hidden_channels=8)
+    for a, b in [
+        (complex_model.node_emb, complex_model.node_emb_im),
+        (complex_model.rel_emb, complex_model.rel_emb_im),
+        (rotate_model.node_emb, rotate_model.node_emb_im),
+    ]:
+        assert not np.allclose(ops.convert_to_numpy(a.embeddings), ops.convert_to_numpy(b.embeddings))