💃 🥤 Extract interaction function from models

docs/source/reference/models.rst

@@ -20,3 +20,15 @@ Extra Modules
 -------------
 .. automodule:: pykeen.nn
     :members:
+
+Interactors
+-----------
+Functional Interface
+~~~~~~~~~~~~~~~~~~~~
+.. automodule:: pykeen.nn.functional
+    :members:
+
+Module Interface
+~~~~~~~~~~~~~~~~
+.. automodule:: pykeen.nn.modules
+    :members:

src/pykeen/models/__init__.py

@@ -6,10 +6,12 @@
 score value is model-dependent, and usually it cannot be directly interpreted as a probability.
 """  # noqa: D205, D400
 
-import inspect
 from typing import Mapping, Set, Type, Union
 
-from .base import EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, MultimodalModel  # noqa:F401
+from .base import (  # noqa:F401
+    EntityEmbeddingModel, EntityRelationEmbeddingModel, Model,
+    MultimodalModel, SimpleVectorEntityRelationEmbeddingModel,
+)
 from .multimodal import ComplExLiteral, DistMultLiteral
 from .unimodal import (
     ComplEx,
@@ -64,10 +66,14 @@
     'get_model_cls',
 ]
 
+_CONCRETE_BASES = {
+    SimpleVectorEntityRelationEmbeddingModel,
+}
+
 
-def _concrete_subclasses(cls):
+def _concrete_subclasses(cls: Type[Model]):
     for subcls in cls.__subclasses__():
-        if not inspect.isabstract(subcls):
+        if not subcls._is_abstract() and subcls not in _CONCRETE_BASES:
             yield subcls
         yield from _concrete_subclasses(subcls)
 

src/pykeen/models/unimodal/complex.py

@@ -4,22 +4,21 @@
 
 from typing import Optional
 
-import torch
 import torch.nn as nn
 
-from ..base import EntityRelationEmbeddingModel
+from ..base import SimpleVectorEntityRelationEmbeddingModel
 from ...losses import Loss, SoftplusLoss
+from ...nn.modules import ComplExInteractionFunction
 from ...regularizers import LpRegularizer, Regularizer
 from ...triples import TriplesFactory
 from ...typing import DeviceHint
-from ...utils import split_complex
 
 __all__ = [
     'ComplEx',
 ]
 
 
-class ComplEx(EntityRelationEmbeddingModel):
+class ComplEx(SimpleVectorEntityRelationEmbeddingModel):
     r"""An implementation of ComplEx [trouillon2016]_.
 
     ComplEx is an extension of :class:`pykeen.models.DistMult` that uses complex valued representations for the
@@ -94,8 +93,11 @@ def __init__(
         :param regularizer: BaseRegularizer
             The regularizer to use.
         """
+        interaction_function = ComplExInteractionFunction()
+
         super().__init__(
             triples_factory=triples_factory,
+            interaction_function=interaction_function,
             embedding_dim=2 * embedding_dim,  # complex embeddings
             automatic_memory_optimization=automatic_memory_optimization,
             loss=loss,
@@ -107,68 +109,3 @@ def __init__(
             entity_initializer=nn.init.normal_,
             relation_initializer=nn.init.normal_,
         )
-
-    @staticmethod
-    def interaction_function(
-        h: torch.FloatTensor,
-        r: torch.FloatTensor,
-        t: torch.FloatTensor,
-    ) -> torch.FloatTensor:
-        """Evaluate the interaction function of ComplEx for given embeddings.
-
-        The embeddings have to be in a broadcastable shape.
-
-        :param h:
-            Head embeddings.
-        :param r:
-            Relation embeddings.
-        :param t:
-            Tail embeddings.
-
-        :return: shape: (...)
-            The scores.
-        """
-        # split into real and imaginary part
-        (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)]
-
-        # ComplEx space bilinear product
-        # *: Elementwise multiplication
-        return sum(
-            (hh * rr * tt).sum(dim=-1)
-            for hh, rr, tt in [
-                (h_re, r_re, t_re),
-                (h_re, r_im, t_im),
-                (h_im, r_re, t_im),
-                (h_im, r_im, t_re),
-            ]
-        )
-
-    def forward(
-        self,
-        h_indices: Optional[torch.LongTensor],
-        r_indices: Optional[torch.LongTensor],
-        t_indices: Optional[torch.LongTensor],
-    ) -> torch.FloatTensor:
-        """Unified score function."""
-        # get embeddings
-        h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices)
-        r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices)
-        t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices)
-
-        # Regularization
-        self.regularize_if_necessary(h, r, t)
-
-        # Compute scores
-        return self.interaction_function(h=h, r=r, t=t)
-
-    def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor:  # noqa: D102
-        return self(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]).view(-1, 1)
-
-    def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor:  # noqa: D102
-        return self(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], t_indices=None)
-
-    def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor:  # noqa: D102
-        return self(h_indices=ht_batch[:, 0], r_indices=None, t_indices=ht_batch[:, 1])
-
-    def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor:  # noqa: D102
-        return self(h_indices=None, r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1])