files: import math
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class ComplExScorer(RelationalScorer):
      r"""Implementation of the ComplEx KGE scorer.
  
      Reference: Théo Trouillon, Johannes Welbl, Sebastian Riedel, Éric Gaussier and
      Guillaume Bouchard: Complex Embeddings for Simple Link Prediction. ICML 2016.
      `<http://proceedings.mlr.press/v48/trouillon16.pdf>`_
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          # Here we use a fast implementation of computing the ComplEx scores using
          # Hadamard products, as in Eq. (11) of paper.
          #
          # Split the relation and object embeddings into real part (first half) and
          # imaginary part (second half).
          p_emb_re, p_emb_im = (t.contiguous() for t in p_emb.chunk(2, dim=1))
          o_emb_re, o_emb_im = (t.contiguous() for t in o_emb.chunk(2, dim=1))
  
          # combine them again to create a column block for each required combination
          s_all = torch.cat((s_emb, s_emb), dim=1)  # re, im, re, im
          r_all = torch.cat((p_emb_re, p_emb, -p_emb_im), dim=1)  # re, re, im, -im
          o_all = torch.cat((o_emb, o_emb_im, o_emb_re), dim=1)  # re, im, im, re
  
          if combine == "spo":
              out = (s_all * o_all * r_all).sum(dim=1)
          elif combine == "sp*":
              out = (s_all * r_all).mm(o_all.transpose(0, 1))
          elif combine == "*po":
              out = (r_all * o_all).mm(s_all.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(n, -1)
  
  
  class ComplEx(KgeModel):
      r"""Implementation of the ComplEx KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(
              config, dataset, ComplExScorer, configuration_key=configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/complex.py
files: import math
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class ComplExScorer(RelationalScorer):
      r"""Implementation of the ComplEx KGE scorer.
  
      Reference: Théo Trouillon, Johannes Welbl, Sebastian Riedel, Éric Gaussier and
      Guillaume Bouchard: Complex Embeddings for Simple Link Prediction. ICML 2016.
      `<http://proceedings.mlr.press/v48/trouillon16.pdf>`_
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          # Here we use a fast implementation of computing the ComplEx scores using
          # Hadamard products, as in Eq. (11) of paper.
          #
          # Split the relation and object embeddings into real part (first half) and
          # imaginary part (second half).
          p_emb_re, p_emb_im = (t.contiguous() for t in p_emb.chunk(2, dim=1))
          o_emb_re, o_emb_im = (t.contiguous() for t in o_emb.chunk(2, dim=1))
  
          # combine them again to create a column block for each required combination
          s_all = torch.cat((s_emb, s_emb), dim=1)  # re, im, re, im
          r_all = torch.cat((p_emb_re, p_emb, -p_emb_im), dim=1)  # re, re, im, -im
          o_all = torch.cat((o_emb, o_emb_im, o_emb_re), dim=1)  # re, im, im, re
  
          if combine == "spo":
              out = (s_all * o_all * r_all).sum(dim=1)
          elif combine == "sp*":
              out = (s_all * r_all).mm(o_all.transpose(0, 1))
          elif combine == "*po":
              out = (r_all * o_all).mm(s_all.transpose(0, 1))
          else:
              raise ValueError('cannot handle combine="{}".format(combine)')
  
          return out.view(n, -1)
  
  
  class ComplEx(KgeModel):
      r"""Implementation of the ComplEx KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(
              config, dataset, ComplExScorer, configuration_key=configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/complex.py
files: import math
  
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class DistMultScorer(RelationalScorer):
      r"""Implementation of the DistMult KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          if combine == "spo":
              out = (s_emb * p_emb * o_emb).sum(dim=1)
          elif combine == "sp*":
              out = (s_emb * p_emb).mm(o_emb.transpose(0, 1))
          elif combine == "*po":
              out = (o_emb * p_emb).mm(s_emb.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(n, -1)
  
  
  class DistMult(KgeModel):
      r"""Implementation of the DistMult KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(
              config, dataset, DistMultScorer, configuration_key=configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/distmult.py
files: import math
  
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class DistMultScorer(RelationalScorer):
      r"""Implementation of the DistMult KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          if combine == "spo":
              out = (s_emb * p_emb * o_emb).sum(dim=1)
          elif combine == "sp*":
              out = (s_emb * p_emb).mm(o_emb.transpose(0, 1))
          elif combine == "*po":
              out = (o_emb * p_emb).mm(s_emb.transpose(0, 1))
          else:
              raise ValueError('cannot handle combine="{}".format(combine)')
  
          return out.view(n, -1)
  
  
  class DistMult(KgeModel):
      r"""Implementation of the DistMult KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(
              config, dataset, DistMultScorer, configuration_key=configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/distmult.py
files: import math
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class FreexScorer(RelationalScorer):
      """Mixing matrix with ComplEx pattern of entries, but entries free.
  
      Initial sandbox model for experimentation with sparsity patterns.
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          o_emb1, o_emb2 = (t.contiguous() for t in o_emb.chunk(2, dim=1))
          s_all = torch.cat((s_emb, s_emb), dim=1)
          o_all = torch.cat((o_emb, o_emb2, o_emb1), dim=1)
  
          if combine == "spo":
              out = (s_all * o_all * p_emb).sum(dim=1)
          elif combine == "sp*":
              out = (s_all * p_emb).mm(o_all.transpose(0, 1))
          elif combine == "*po":
              out = (p_emb * o_all).mm(s_all.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(n, -1)
  
  
  class Freex(KgeModel):
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
  
          # set relation embedder dimensionality
          if self.get_option("relation_embedder.dim") < 0:
              self.config.set(
                  self.configuration_key + ".relation_embedder.dim",
                  2 * self.get_option("entity_embedder.dim"),
                  log=True,
              )
  
          super().__init__(
              config, dataset, FreexScorer, configuration_key=self.configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/experimental/freex.py
files: import math
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class FreexScorer(RelationalScorer):
      """Mixing matrix with ComplEx pattern of entries, but entries free.
  
      Initial sandbox model for experimentation with sparsity patterns.
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          o_emb1, o_emb2 = (t.contiguous() for t in o_emb.chunk(2, dim=1))
          s_all = torch.cat((s_emb, s_emb), dim=1)
          o_all = torch.cat((o_emb, o_emb2, o_emb1), dim=1)
  
          if combine == "spo":
              out = (s_all * o_all * p_emb).sum(dim=1)
          elif combine == "sp*":
              out = (s_all * p_emb).mm(o_all.transpose(0, 1))
          elif combine == "*po":
              out = (p_emb * o_all).mm(s_all.transpose(0, 1))
          else:
              raise ValueError('cannot handle combine="{}".format(combine)')
  
          return out.view(n, -1)
  
  
  class Freex(KgeModel):
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
  
          # set relation embedder dimensionality
          if self.get_option("relation_embedder.dim") < 0:
              self.config.set(
                  self.configuration_key + ".relation_embedder.dim",
                  2 * self.get_option("entity_embedder.dim"),
                  log=True,
              )
  
          super().__init__(
              config, dataset, FreexScorer, configuration_key=self.configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/experimental/freex.py
files: import importlib
  import tempfile
  
  from torch import Tensor
  import torch.nn
  
  import kge
  from kge import Config, Configurable, Dataset
  from kge.misc import filename_in_module
  from typing import Any, Dict, List, Optional, Union
  
  from typing import TYPE_CHECKING
  
  if TYPE_CHECKING:
      from kge.job import Job
  
  
  SLOTS = [0, 1, 2]
  S, P, O = SLOTS
  
  
  class KgeBase(torch.nn.Module, Configurable):
      r"""Base class for all KGE models, scorers, and embedders."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          Configurable.__init__(self, config, configuration_key)
          torch.nn.Module.__init__(self)
          self.dataset = dataset
          self.meta: Dict[str, Any] = dict()  #: meta-data stored with this module
  
      def initialize(self, what: Tensor, initialize: str, initialize_args):
          try:
              getattr(torch.nn.init, initialize)(what, **initialize_args)
          except:
              raise ValueError(
                  "invalid initialization options: {} with args {}".format(
                      initialize, initialize_args
                  )
              )
  
      def prepare_job(self, job: "Job", **kwargs):
          r"""Prepares the given job to work with this model.
  
          If this model does not support the specified job type, this function may raise
          an error.
  
          This function commonly registers hooks specific to this model. For a list of
          available hooks during training or evaluation, see :class:`TrainingJob` or
          :class:`EvaluationJob`:, respectively.
  
          """
  
      def penalty(self, **kwargs) -> List[Tensor]:
          r"""Returns additional penalty terms that are added to the loss during training.
  
          This method is called once per batch during training. The arguments being passed
          depend on the trainer being used.
  
          Returns a (possibly empty) list of penalty terms.
  
          """
  
          return []
  
      def save(self):
          "Returns data structure to save state"
          return (self.state_dict(), self.meta)
  
      def load(self, savepoint):
          "Loads state from a saved data structure"
          self.load_state_dict(savepoint[0])
          self.meta = savepoint[1]
  
  
  class RelationalScorer(KgeBase):
      r"""Base class for all relational scorers.
  
      Relational scorers take as input the embeddings of (subject, predicate,
      object)-triple and produce a score.
  
      Implementations of this class should either implement
      :func:`~RelationalScorer.score_emb_spo` (the quick way, but potentially inefficient)
      or :func:`~RelationalScorer.score_emb` (the hard way, potentially more efficient).
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key: str):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb_spo(self, s_emb: Tensor, p_emb: Tensor, o_emb: Tensor) -> Tensor:
          r"""Scores a set of triples specified by their embeddings.
  
          `s_emb`, `p_emb`, and `o_emb` are tensors of size :math:`n\times d_e`,
          :math:`n\times d_r`, and :math:`n\times d_e`, where :math:`d_e` and
          :math:`d_r` are the sizes of the entity and relation embeddings, respectively.
  
          The embeddings are combined row-wise. The output is a :math`n\times 1` tensor,
          in which the :math:`i`-th entry holds the score of the embedding triple
          :math:`(s_i, p_i, o_i)`.
  
          """
          return self.score_emb(s_emb, p_emb, o_emb, "spo")
  
      def score_emb(
          self, s_emb: Tensor, p_emb: Tensor, o_emb: Tensor, combine: str
      ) -> Tensor:
          r"""Scores a set of triples specified by their embeddings.
  
          `s_emb`, `p_emb`, and `o_emb` are tensors of size :math:`n_s\times d_e`,
          :math:`n_p\times d_r`, and :math:`n_o\times d_e`, where :math:`d_e` and
          :math:`d_r` are the sizes of the entity and relation embeddings, respectively.
  
          The provided embeddings are combined based on the value of `combine`. Common
          values are :code:`"spo"`, :code:`"sp*"`, and :code:`"*po"`. Not all models may
          support all combinations.
  
          When `combine` is :code:`"spo"`, then embeddings are combined row-wise. In this
          case, it is required that :math:`n_s=n_p=n_o=n`. The output is identical to
          :func:`~RelationalScorer.score_emb_spo`, i.e., a :math`n\times 1` tensor, in
          which the :math:`i`-th entry holds the score of the embedding triple
          :math:`(s_i, p_i, o_i)`.
  
          When `combine` is :code:`"
  
  [... the rest is omitted ...]
files: kge/model/kge_model.py
files: import importlib
  import tempfile
  
  from torch import Tensor
  import torch.nn
  
  import kge
  from kge import Config, Configurable, Dataset
  from kge.misc import filename_in_module
  from typing import Any, Dict, List, Optional, Union
  
  from typing import TYPE_CHECKING
  
  if TYPE_CHECKING:
      from kge.job import Job
  
  
  SLOTS = [0, 1, 2]
  S, P, O = SLOTS
  
  
  class KgeBase(torch.nn.Module, Configurable):
      r"""Base class for all KGE models, scorers, and embedders."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          Configurable.__init__(self, config, configuration_key)
          torch.nn.Module.__init__(self)
          self.dataset = dataset
          self.meta: Dict[str, Any] = dict()  #: meta-data stored with this module
  
      def initialize(self, what: Tensor, initialize: str, initialize_args):
          try:
              getattr(torch.nn.init, initialize)(what, **initialize_args)
          except:
              raise ValueError(
                  "invalid initialization options: {} with args {}".format(
                      initialize, initialize_args
                  )
              )
  
      def prepare_job(self, job: "Job", **kwargs):
          r"""Prepares the given job to work with this model.
  
          If this model does not support the specified job type, this function may raise
          an error.
  
          This function commonly registers hooks specific to this model. For a list of
          available hooks during training or evaluation, see :class:`TrainingJob` or
          :class:`EvaluationJob`:, respectively.
  
          """
  
      def penalty(self, **kwargs) -> List[Tensor]:
          r"""Returns additional penalty terms that are added to the loss during training.
  
          This method is called once per batch during training. The arguments being passed
          depend on the trainer being used.
  
          Returns a (possibly empty) list of penalty terms.
  
          """
  
          return []
  
      def save(self):
          "Returns data structure to save state"
          return (self.state_dict(), self.meta)
  
      def load(self, savepoint):
          "Loads state from a saved data structure"
          self.load_state_dict(savepoint[0])
          self.meta = savepoint[1]
  
  
  class RelationalScorer(KgeBase):
      r"""Base class for all relational scorers.
  
      Relational scorers take as input the embeddings of (subject, predicate,
      object)-triple and produce a score.
  
      Implementations of this class should either implement
      :func:`~RelationalScorer.score_emb_spo` (the quick way, but potentially inefficient)
      or :func:`~RelationalScorer.score_emb` (the hard way, potentially more efficient).
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key: str):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb_spo(self, s_emb: Tensor, p_emb: Tensor, o_emb: Tensor) -> Tensor:
          r"""Scores a set of triples specified by their embeddings.
  
          `s_emb`, `p_emb`, and `o_emb` are tensors of size :math:`n\times d_e`,
          :math:`n\times d_r`, and :math:`n\times d_e`, where :math:`d_e` and
          :math:`d_r` are the sizes of the entity and relation embeddings, respectively.
  
          The embeddings are combined row-wise. The output is a :math`n\times 1` tensor,
          in which the :math:`i`-th entry holds the score of the embedding triple
          :math:`(s_i, p_i, o_i)`.
  
          """
          return self.score_emb(s_emb, p_emb, o_emb, "spo")
  
      def score_emb(
          self, s_emb: Tensor, p_emb: Tensor, o_emb: Tensor, combine: str
      ) -> Tensor:
          r"""Scores a set of triples specified by their embeddings.
  
          `s_emb`, `p_emb`, and `o_emb` are tensors of size :math:`n_s\times d_e`,
          :math:`n_p\times d_r`, and :math:`n_o\times d_e`, where :math:`d_e` and
          :math:`d_r` are the sizes of the entity and relation embeddings, respectively.
  
          The provided embeddings are combined based on the value of `combine`. Common
          values are :code:`"spo"`, :code:`"sp*"`, and :code:`"*po"`. Not all models may
          support all combinations.
  
          When `combine` is :code:`"spo"`, then embeddings are combined row-wise. In this
          case, it is required that :math:`n_s=n_p=n_o=n`. The output is identical to
          :func:`~RelationalScorer.score_emb_spo`, i.e., a :math`n\times 1` tensor, in
          which the :math:`i`-th entry holds the score of the embedding triple
          :math:`(s_i, p_i, o_i)`.
  
          When `combine` is :code:`"
  
  [... the rest is omitted ...]
files: kge/model/kge_model.py
files: import torch
  from torch import Tensor
  from kge import Config, Dataset
  from kge.model.kge_model import KgeModel
  
  
  class ReciprocalRelationsModel(KgeModel):
      """Modifies a base model to use different relation embeddings for predicting subject and object.
  
      This implements the reciprocal relations training procedure of [TODO cite ConvE].
      Note that this model cannot be used to score a single triple, but only to rank sp*
      or *po questions.
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
  
          # Initialize base model
          # Using a dataset with twice the number of relations to initialize base model
          alt_dataset = Dataset(
              dataset.config,
              dataset.num_entities,
              dataset.entities,
              dataset.num_relations * 2,
              dataset.relations,
              dataset.train,
              dataset.train_meta,
              dataset.valid,
              dataset.valid_meta,
              dataset.test,
              dataset.test_meta,
          )
          base_model = KgeModel.create(
              config, alt_dataset, self.configuration_key + ".base_model"
          )
  
          # Initialize this model
          super().__init__(
              config, dataset, base_model.get_scorer(), initialize_embedders=False
          )
          self._base_model = base_model
          # TODO change entity_embedder assignment to sub and obj embedders when support
          # for that is added
          self._entity_embedder = self._base_model.get_s_embedder()
          self._relation_embedder = self._base_model.get_p_embedder()
  
      def prepare_job(self, job, **kwargs):
          self._base_model.prepare_job(job, **kwargs)
  
      def penalty(self, **kwargs):
          return super().penalty(**kwargs) + self._base_model.penalty(**kwargs)
  
      def score_spo(self, s: Tensor, p: Tensor, o: Tensor, direction=None) -> Tensor:
          if direction == "o":
              return super().score_spo(s, p, o, "o")
          elif direction == "s":
              return super().score_spo(o, p + self.dataset.num_relations, s, "o")
          else:
              raise Exception(
                  "The reciprocal relations model cannot compute "
                  "undirected spo scores."
              )
  
      def score_po(self, p, o, s=None):
          if s is None:
              s = self.get_s_embedder().embed_all()
          else:
              s = self.get_s_embedder().embed(s)
          p = self.get_p_embedder().embed(p + self.dataset.num_relations)
          o = self.get_o_embedder().embed(o)
          return self._scorer.score_emb(o, p, s, combine="sp*")
  
      def score_so(self, s, o, p=None):
          raise Exception("The reciprocal relations model cannot score relations.")
  
      def score_sp_po(
          self,
          s: torch.Tensor,
          p: torch.Tensor,
          o: torch.Tensor,
          entity_subset: torch.Tensor = None,
      ) -> torch.Tensor:
          s = self.get_s_embedder().embed(s)
          p_inv = self.get_p_embedder().embed(p + self.dataset.num_relations)
          p = self.get_p_embedder().embed(p)
          o = self.get_o_embedder().embed(o)
          if self.get_s_embedder() is self.get_o_embedder():
              if entity_subset is not None:
                  all_entities = self.get_s_embedder().embed(entity_subset)
              else:
                  all_entities = self.get_s_embedder().embed_all()
              sp_scores = self._scorer.score_emb(s, p, all_entities, combine="sp*")
              po_scores = self._scorer.score_emb(o, p_inv, all_entities, combine="sp*")
          else:
              if entity_subset is not None:
                  all_objects = self.get_o_embedder().embed(entity_subset)
                  all_subjects = self.get_s_embedder().embed(entity_subset)
              else:
                  all_objects = self.get_o_embedder().embed_all()
                  all_subjects = self.get_s_embedder().embed_all()
              sp_scores = self._scorer.score_emb(s, p, all_objects, combine="sp*")
              po_scores = self._scorer.score_emb(o, p_inv, all_subjects, combine="sp*")
          return torch.cat((sp_scores, po_scores), dim=1)
  
  
  [... the rest is omitted ...]
files: kge/model/reciprocal_relations_model.py
files: import torch
  from torch import Tensor
  from kge import Config, Dataset
  from kge.model.kge_model import KgeModel
  
  
  class ReciprocalRelationsModel(KgeModel):
      """Modifies a base model to use different relation embeddings for predicting subject and object.
  
      This implements the reciprocal relations training procedure of [TODO cite ConvE].
      Note that this model cannot be used to score a single triple, but only to rank sp*
      or *po questions.
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
  
          # Initialize base model
          # Using a dataset with twice the number of relations to initialize base model
          alt_dataset = Dataset(
              dataset.config,
              dataset.num_entities,
              dataset.entities,
              dataset.num_relations * 2,
              dataset.relations,
              dataset.train,
              dataset.train_meta,
              dataset.valid,
              dataset.valid_meta,
              dataset.test,
              dataset.test_meta,
          )
          base_model = KgeModel.create(
              config, alt_dataset, self.configuration_key + ".base_model"
          )
  
          # Initialize this model
          super().__init__(
              config, dataset, base_model.get_scorer(), initialize_embedders=False
          )
          self._base_model = base_model
          # TODO change entity_embedder assignment to sub and obj embedders when support
          # for that is added
          self._entity_embedder = self._base_model.get_s_embedder()
          self._relation_embedder = self._base_model.get_p_embedder()
  
      def prepare_job(self, job, **kwargs):
          self._base_model.prepare_job(job, **kwargs)
  
      def penalty(self, **kwargs):
          return super().penalty(**kwargs) + self._base_model.penalty(**kwargs)
  
      def score_spo(self, s: Tensor, p: Tensor, o: Tensor, direction=None) -> Tensor:
          if direction == "o":
              return super().score_spo(s, p, o, "o")
          elif direction == "s":
              return super().score_spo(o, p + self.dataset.num_relations, s, "o")
          else:
              raise Exception(
                  "The reciprocal relations model cannot compute "
                  "undirected spo scores."
              )
  
      def score_po(self, p, o, s=None):
          if s is None:
              s = self.get_s_embedder().embed_all()
          else:
              s = self.get_s_embedder().embed(s)
          p = self.get_p_embedder().embed(p + self.dataset.num_relations)
          o = self.get_o_embedder().embed(o)
          return self._scorer.score_emb(o, p, s, combine="sp*")
  
      def score_sp_po(self, s: torch.Tensor, p: torch.Tensor, o: torch.Tensor,
                      entity_subset: torch.Tensor = None) -> torch.Tensor:
          s = self.get_s_embedder().embed(s)
          p_inv = self.get_p_embedder().embed(p + self.dataset.num_relations)
          p = self.get_p_embedder().embed(p)
          o = self.get_o_embedder().embed(o)
          if self.get_s_embedder() is self.get_o_embedder():
              if entity_subset is not None:
                  all_entities = self.get_s_embedder().embed(entity_subset)
              else:
                  all_entities = self.get_s_embedder().embed_all()
              sp_scores = self._scorer.score_emb(s, p, all_entities, combine="sp*")
              po_scores = self._scorer.score_emb(o, p_inv, all_entities, combine="sp*")
          else:
              if entity_subset is not None:
                  all_objects = self.get_o_embedder().embed(entity_subset)
                  all_subjects = self.get_s_embedder().embed(entity_subset)
              else:
                  all_objects = self.get_o_embedder().embed_all()
                  all_subjects = self.get_s_embedder().embed_all()
              sp_scores = self._scorer.score_emb(s, p, all_objects, combine="sp*")
              po_scores = self._scorer.score_emb(o, p_inv, all_subjects, combine="sp*")
          return torch.cat((sp_scores, po_scores), dim=1)
  
  
  [... the rest is omitted ...]
files: kge/model/reciprocal_relations_model.py
files: import math
  
  import torch
  
  from kge import Config, Dataset
  from kge.model.kge_model import KgeEmbedder, KgeModel, RelationalScorer
  
  
  class RescalScorer(RelationalScorer):
      r"""Implementation of the RESCAL KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(
          self,
          s_emb: torch.Tensor,
          p_emb: torch.Tensor,
          o_emb: torch.Tensor,
          combine: str,
      ):
          batch_size = p_emb.size(0)
          entity_size = s_emb.size(-1)
  
          # reshape relation embeddings to obtain mixing matrices for RESCAL
          p_mixmat = p_emb.view(-1, entity_size, entity_size)
  
          if combine == "spo":
              out = (
                  s_emb.unsqueeze(1)  # [batch x 1 x entity_size]
                  .bmm(p_mixmat)  # apply mixing matrices
                  .view(batch_size, entity_size)  # drop dim 1
                  * o_emb  # apply object embeddings
              ).sum(
                  dim=-1
              )  # and sum to obtain predictions
          elif combine == "sp*":
              out = (
                  s_emb.unsqueeze(1)
                  .bmm(p_mixmat)
                  .view(batch_size, entity_size)
                  .mm(o_emb.transpose(0, 1))
              )
          elif combine == "*po":
              out = (
                  p_mixmat.bmm(o_emb.unsqueeze(2))
                  .view(batch_size, entity_size)
                  .mm(s_emb.transpose(0, 1))
              )
          else:
              super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(batch_size, -1)
  
  
  class Rescal(KgeModel):
      r"""Implementation of the RÃ‰SCAL KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          rescal_set_relation_embedder_dim(
              config, dataset, self.configuration_key + ".relation_embedder"
          )
          super().__init__(
              config, dataset, RescalScorer, configuration_key=self.configuration_key
          )
  
  
  def rescal_set_relation_embedder_dim(config, dataset, rel_emb_conf_key):
      """Set the relation embedder dimensionality for RESCAL in the config.
  
      If <0, set it to the square of the size of the entity embedder. Else leave
      unchanged.
  
      """
      dim = config.get_default(rel_emb_conf_key + ".dim")
      if dim < 0:  # autodetect relation embedding dimensionality
          ent_emb_conf_key = rel_emb_conf_key.replace(
              "relation_embedder", "entity_embedder"
          )
          if ent_emb_conf_key == rel_emb_conf_key:
              raise ValueError(
                  "Cannot determine relation embedding size; please set manually."
              )
          dim = config.get_default(ent_emb_conf_key + ".dim") ** 2
          config.set(rel_emb_conf_key + ".dim", dim, log=True)
  
  
  [... the rest is omitted ...]
files: kge/model/rescal.py
files: import math
  
  import torch
  
  from kge import Config, Dataset
  from kge.model.kge_model import KgeEmbedder, KgeModel, RelationalScorer
  
  
  class RescalScorer(RelationalScorer):
      r"""Implementation of the RESCAL KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(
          self,
          s_emb: torch.Tensor,
          p_emb: torch.Tensor,
          o_emb: torch.Tensor,
          combine: str,
      ):
          batch_size = p_emb.size(0)
          entity_size = s_emb.size(-1)
  
          # reshape relation embeddings to obtain mixing matrices for RESCAL
          p_mixmat = p_emb.view(-1, entity_size, entity_size)
  
          if combine == "spo":
              out = (
                  s_emb.unsqueeze(1)  # [batch x 1 x entity_size]
                  .bmm(p_mixmat)  # apply mixing matrices
                  .view(batch_size, entity_size)  # drop dim 1
                  * o_emb  # apply object embeddings
              ).sum(
                  dim=-1
              )  # and sum to obtain predictions
          elif combine == "sp*":
              out = (
                  s_emb.unsqueeze(1)
                  .bmm(p_mixmat)
                  .view(batch_size, entity_size)
                  .mm(o_emb.transpose(0, 1))
              )
          elif combine == "*po":
              out = (
                  p_mixmat.bmm(o_emb.unsqueeze(2))
                  .view(batch_size, entity_size)
                  .mm(s_emb.transpose(0, 1))
              )
          else:
              raise ValueError('cannot handle combine="{}".format(combine)')
  
          return out.view(batch_size, -1)
  
  
  class Rescal(KgeModel):
      r"""Implementation of the RÃ‰SCAL KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          rescal_set_relation_embedder_dim(
              config, dataset, self.configuration_key + ".relation_embedder"
          )
          super().__init__(
              config, dataset, RescalScorer, configuration_key=self.configuration_key
          )
  
  
  def rescal_set_relation_embedder_dim(config, dataset, rel_emb_conf_key):
      """Set the relation embedder dimensionality for RESCAL in the config.
  
      If <0, set it to the square of the size of the entity embedder. Else leave
      unchanged.
  
      """
      dim = config.get_default(rel_emb_conf_key + ".dim")
      if dim < 0:  # autodetect relation embedding dimensionality
          ent_emb_conf_key = rel_emb_conf_key.replace(
              "relation_embedder", "entity_embedder"
          )
          if ent_emb_conf_key == rel_emb_conf_key:
              raise ValueError(
                  "Cannot determine relation embedding size; please set manually."
              )
          dim = config.get_default(ent_emb_conf_key + ".dim") ** 2
          config.set(rel_emb_conf_key + ".dim", dim, log=True)
  
  
  [... the rest is omitted ...]
files: kge/model/rescal.py
files: import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  from torch.nn import functional as F
  
  
  class TransEScorer(RelationalScorer):
      r"""Implementation of the TransE KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
          self._norm = self.get_option("l_norm")
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
          if combine == "spo":
              out = -F.pairwise_distance(s_emb + p_emb, o_emb, p=self._norm)
          elif combine == "sp*":
              out = -torch.cdist(s_emb + p_emb, o_emb, p=self._norm)
              # old pytorch 1.2 version
              # sp_emb = s_emb + p_emb
              # out = torch.zeros(n, o_emb.size(0)).to(self.config.get("job.device"))
              # for i in range(n):
              #     out[i, :] = -F.pairwise_distance(sp_emb[i, :], o_emb, p=self._norm)
          elif combine == "*po":
              out = -torch.cdist(o_emb - p_emb, s_emb, p=self._norm)
              # old pytorch 1.2 version
              # po_emb = o_emb - p_emb
              # out = torch.zeros(n, s_emb.size(0)).to(self.config.get("job.device"))
              # for i in range(n):
              #     out[i, :] = -F.pairwise_distance(po_emb[i, :], s_emb, p=self._norm)
          else:
              super().score_emb(s_emb, p_emb, o_emb, combine)
          return out.view(n, -1)
  
  
  class TransE(KgeModel):
      r"""Implementation of the TransE KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(
              config, dataset, TransEScorer, configuration_key=configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/transe.py
files: import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  from torch.nn import functional as F
  
  
  class TransEScorer(RelationalScorer):
      r"""Implementation of the TransE KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
          self._norm = self.get_option("l_norm")
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
          if combine == "spo":
              out = -F.pairwise_distance(s_emb + p_emb, o_emb, p=self._norm)
          elif combine == "sp*":
              out = -torch.cdist(s_emb + p_emb, o_emb, p=self._norm)
              # old pytorch 1.2 version
              # sp_emb = s_emb + p_emb
              # out = torch.zeros(n, o_emb.size(0)).to(self.config.get("job.device"))
              # for i in range(n):
              #     out[i, :] = -F.pairwise_distance(sp_emb[i, :], o_emb, p=self._norm)
          elif combine == "*po":
              out = -torch.cdist(o_emb - p_emb, s_emb, p=self._norm)
              # old pytorch 1.2 version
              # po_emb = o_emb - p_emb
              # out = torch.zeros(n, s_emb.size(0)).to(self.config.get("job.device"))
              # for i in range(n):
              #     out[i, :] = -F.pairwise_distance(po_emb[i, :], s_emb, p=self._norm)
          else:
              raise ValueError('cannot handle combine="{}".format(combine)')
          return out.view(n, -1)
  
  
  class TransE(KgeModel):
      r"""Implementation of the TransE KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(
              config, dataset, TransEScorer, configuration_key=configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/transe.py