files: import itertools
  import os
  import math
  import time
  from collections import defaultdict
  
  from dataclasses import dataclass
  
  import torch
  import torch.utils.data
  
  from kge import Config, Dataset
  from kge.job import Job
  from kge.model import KgeModel
  
  from kge.util import KgeLoss, KgeOptimizer, KgeSampler, KgeLRScheduler
  from typing import Any, Callable, Dict, List, Optional
  import kge.job.util
  
  SLOTS = [0, 1, 2]
  S, P, O = SLOTS
  
  
  class TrainingJob(Job):
      """Abstract base job to train a single model with a fixed set of hyperparameters.
  
      Also used by jobs such as :class:`SearchJob`.
  
      Subclasses for specific training methods need to implement `_prepare` and
      `_process_batch`.
  
      """
  
      def __init__(
          self, config: Config, dataset: Dataset, parent_job: Job = None
      ) -> None:
          from kge.job import EvaluationJob
  
          super().__init__(config, dataset, parent_job)
          self.model: KgeModel = KgeModel.create(config, dataset)
          self.optimizer = KgeOptimizer.create(config, self.model)
          self.kge_lr_scheduler = KgeLRScheduler(config, self.optimizer)
          self.loss = KgeLoss.create(config)
          self.abort_on_nan: bool = config.get("train.abort_on_nan")
          self.batch_size: int = config.get("train.batch_size")
          self.device: str = self.config.get("job.device")
          self.train_split = config
  
  [... the rest is omitted ...]
files: kge/job/train.py
files: import itertools
  import os
  import math
  import time
  from collections import defaultdict
  
  from dataclasses import dataclass
  
  import torch
  import torch.utils.data
  
  from kge import Config, Dataset
  from kge.job import Job
  from kge.model import KgeModel
  
  from kge.util import KgeLoss, KgeOptimizer, KgeSampler, KgeLRScheduler
  from typing import Any, Callable, Dict, List, Optional
  import kge.job.util
  
  SLOTS = [0, 1, 2]
  S, P, O = SLOTS
  
  
  class TrainingJob(Job):
      """Abstract base job to train a single model with a fixed set of hyperparameters.
  
      Also used by jobs such as :class:`SearchJob`.
  
      Subclasses for specific training methods need to implement `_prepare` and
      `_process_batch`.
  
      """
  
      def __init__(
          self, config: Config, dataset: Dataset, parent_job: Job = None
      ) -> None:
          from kge.job import EvaluationJob
  
          super().__init__(config, dataset, parent_job)
          self.model: KgeModel = KgeModel.create(config, dataset)
          self.optimizer = KgeOptimizer.create(config, self.model)
          self.kge_lr_scheduler = KgeLRScheduler(config, self.optimizer)
          self.loss = KgeLoss.create(config)
          self.abort_on_nan: bool = config.get("train.abort_on_nan")
          self.batch_size: int = config.get("train.batch_size")
          self.device: str = self.config.get("job.device")
          self.train_split = config
  
  [... the rest is omitted ...]
files: kge/job/train.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
  
  [... 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
  
  [... the rest is omitted ...]
files: kge/model/complex.py
files: import torch
  import math
  
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class ConvEScorer(RelationalScorer):
      r"""Implementation of the ConvE KGE scorer.
  
      Must be used with ReciprocalRelationsModel."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
          # self.configuration_key = configuration_key
          self.emb_dim = self.get_option("entity_embedder.dim") - 1
          aspect_ratio = self.get_option("2D_aspect_ratio")
          self.emb_height = math.sqrt(self.emb_dim / aspect_ratio)
          self.emb_width = self.emb_height * aspect_ratio
  
          # round embedding dimension to match aspect ratio
          rounded_height = math.ceil(self.emb_height)
          if self.get_option("round_dim") and rounded_height != self.emb_height:
              self.emb_height = rounded_height
              self.emb_width = self.emb_height * aspect_ratio
              self.emb_dim = self.emb_height * self.emb_width
              self.set_option("entity_embedder.dim", self.emb_dim + 1)
              self.set_option("relation_embedder.dim", self.emb_dim + 1)
              config.log(
                  "Rounded embedding dimension up to {} to match given aspect ratio.".format(
                      self.emb_dim
                  )
              )
          elif self.emb_dim % self.emb_height or self.emb_dim % self.emb_width:
              raise Exception(
                  "Aspect ratio {} does not produce 2
  
  [... the rest is omitted ...]
files: kge/model/conve.py
files: import torch
  import math
  
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class ConvEScorer(RelationalScorer):
      r"""Implementation of the ConvE KGE scorer.
  
      Must be used with ReciprocalRelationsModel."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
          # self.configuration_key = configuration_key
          self.emb_dim = self.get_option("entity_embedder.dim") - 1
          aspect_ratio = self.get_option("2D_aspect_ratio")
          self.emb_height = math.sqrt(self.emb_dim / aspect_ratio)
          self.emb_width = self.emb_height * aspect_ratio
  
          # round embedding dimension to match aspect ratio
          rounded_height = math.ceil(self.emb_height)
          if self.get_option("round_dim") and rounded_height != self.emb_height:
              self.emb_height = rounded_height
              self.emb_width = self.emb_height * aspect_ratio
              self.emb_dim = self.emb_height * self.emb_width
              self.set_option("entity_embedder.dim", self.emb_dim + 1)
              self.set_option("relation_embedder.dim", self.emb_dim + 1)
              config.log(
                  "Rounded embedding dimension up to {} to match given aspect ratio.".format(
                      self.emb_dim
                  )
              )
          elif self.emb_dim % self.emb_height or self.emb_dim % self.emb_width:
              raise Exception(
                  "Aspect ratio {} does not produce 2
  
  [... the rest is omitted ...]
files: kge/model/conve.py
files: import math
  
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class CPScorer(RelationalScorer):
      r"""Implementation of the CP 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)
  
          # use only first half for subjects and second half for objects
          half_dim = s_emb.shape[1] // 2
          s_emb_h = s_emb[:, :half_dim]
          o_emb_t = o_emb[:, half_dim:]
  
          if combine == "spo":
              out = (s_emb_h * p_emb * o_emb_t).sum(dim=1)
          elif combine == "sp_":
              out = (s_emb_h * p_emb).mm(o_emb_t.transpose(0, 1))
          elif combine == "_po":
              out = (o_emb_t * p_emb).mm(s_emb_h.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(n, -1)
  
  
  class CP(KgeModel):
      r"""Implementation of the CP KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          if self.get_option("entity_embed
  
  [... the rest is omitted ...]
files: kge/model/cp.py
files: import math
  
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class CPScorer(RelationalScorer):
      r"""Implementation of the CP 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)
  
          # use only first half for subjects and second half for objects
          half_dim = s_emb.shape[1] // 2
          s_emb_h = s_emb[:, :half_dim]
          o_emb_t = o_emb[:, half_dim:]
  
          if combine == "spo":
              out = (s_emb_h * p_emb * o_emb_t).sum(dim=1)
          elif combine == "sp*":
              out = (s_emb_h * p_emb).mm(o_emb_t.transpose(0, 1))
          elif combine == "*po":
              out = (o_emb_t * p_emb).mm(s_emb_h.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(n, -1)
  
  
  class CP(KgeModel):
      r"""Implementation of the CP KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          if self.get_option("entity_embed
  
  [... the rest is omitted ...]
files: kge/model/cp.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:
              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 importlib
  import tempfile
  
  from torch import Tensor
  import torch.nn
  import os
  
  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.
  
         
  
  [... the rest is omitted ...]
files: kge/model/kge_model.py
files: import importlib
  import tempfile
  
  from torch import Tensor
  import torch.nn
  import os
  
  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.
  
         
  
  [... 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.shallow_copy()
          alt_dataset._num_relations = dataset.num_relations() * 2
          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_sp
  
  [... 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.shallow_copy()
          alt_dataset._num_relations = dataset.num_relations() * 2
          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_sp
  
  [... 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
  
  [... 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
  
  [... 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
  
  
  # TODO sp_ and _po scoring with RotatE leads to *large* intermediate results. It's
  # unclear whether this can be fixed. Expect out-of-memory errors when using RotatE with
  # 1vsAll or KvsAll training. To do validation/evaluation, you may want to set
  # eval.chunk_size.
  class RotatEScorer(RelationalScorer):
      r"""Implementation of the RotatE 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)
  
          # determine real and imaginary part
          s_emb_re, s_emb_im = torch.chunk(s_emb, 2, dim=1)
          o_emb_re, o_emb_im = torch.chunk(o_emb, 2, dim=1)
  
          # convert from radians to points on complex unix ball
          p_emb_re, p_emb_im = torch.cos(p_emb), torch.sin(p_emb)
  
          if combine == "spo":
              # compute the difference vector (s*p-t)
              sp_emb_re, sp_emb_im = hadamard_complex(
                  s_emb_re, s_emb_im, p_emb_re,
  
  [... the rest is omitted ...]
files: kge/model/rotate.py
files: import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  from torch.nn import functional as F
  
  
  # TODO sp* and *po scoring with RotatE leads to *large* intermediate results. It's
  # unclear whether this can be fixed. Expect out-of-memory errors when using RotatE with
  # 1vsAll or KvsAll training. To do validation/evaluation, you may want to set
  # eval.chunk_size.
  class RotatEScorer(RelationalScorer):
      r"""Implementation of the RotatE 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)
  
          # determine real and imaginary part
          s_emb_re, s_emb_im = torch.chunk(s_emb, 2, dim=1)
          o_emb_re, o_emb_im = torch.chunk(o_emb, 2, dim=1)
  
          # convert from radians to points on complex unix ball
          p_emb_re, p_emb_im = torch.cos(p_emb), torch.sin(p_emb)
  
          if combine == "spo":
              # compute the difference vector (s*p-t)
              sp_emb_re, sp_emb_im = hadamard_complex(
                  s_emb_re, s_emb_im, p_emb_re,
  
  [... the rest is omitted ...]
files: kge/model/rotate.py
files: import math
  
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class SimplEScorer(RelationalScorer):
      r"""Implementation of the SimplE 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)
  
          # split left/right
          s_emb_h, s_emb_t = torch.chunk(s_emb, 2, dim=1)
          p_emb_forward, p_emb_backward = torch.chunk(p_emb, 2, dim=1)
          o_emb_h, o_emb_t = torch.chunk(o_emb, 2, dim=1)
  
          if combine == "spo":
              out1 = (s_emb_h * p_emb_forward * o_emb_t).sum(dim=1)
              out2 = (s_emb_t * p_emb_backward * o_emb_h).sum(dim=1)
          elif combine == "sp_":
              out1 = (s_emb_h * p_emb_forward).mm(o_emb_t.transpose(0, 1))
              out2 = (s_emb_t * p_emb_backward).mm(o_emb_h.transpose(0, 1))
          elif combine == "_po":
              out1 = (o_emb_t * p_emb_forward).mm(s_emb_h.transpose(0
  
  [... the rest is omitted ...]
files: kge/model/simple.py
files: import math
  
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class SimplEScorer(RelationalScorer):
      r"""Implementation of the SimplE 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)
  
          # split left/right
          s_emb_h, s_emb_t = torch.chunk(s_emb, 2, dim=1)
          p_emb_forward, p_emb_backward = torch.chunk(p_emb, 2, dim=1)
          o_emb_h, o_emb_t = torch.chunk(o_emb, 2, dim=1)
  
          if combine == "spo":
              out1 = (s_emb_h * p_emb_forward * o_emb_t).sum(dim=1)
              out2 = (s_emb_t * p_emb_backward * o_emb_h).sum(dim=1)
          elif combine == "sp*":
              out1 = (s_emb_h * p_emb_forward).mm(o_emb_t.transpose(0, 1))
              out2 = (s_emb_t * p_emb_backward).mm(o_emb_h.transpose(0, 1))
          elif combine == "*po":
              out1 = (o_emb_t * p_emb_forward).mm(s_emb_h.transpose(0
  
  [... the rest is omitted ...]
files: kge/model/simple.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)
          elif combine == "_po":
              out = -torch.cdist(o_emb - p_emb, 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)
          elif combine == "*po":
              out = -torch.cdist(o_emb - p_emb, 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