class_docstr: DGL Knowledge Embedding Model.
  
  Parameters
  ----------
  args:
      Global configs.
  model_name : str
      Which KG model to use, including 'TransE_l1', 'TransE_l2', 'TransR',
      'RESCAL', 'DistMult', 'ComplEx', 'RotatE', 'SimplE'
  n_entities : int
      Num of entities.
  n_relations : int
      Num of relations.
  hidden_dim : int
      Dimension size of embedding.
  gamma : float
      Gamma for score function.
  double_entity_emb : bool
      If True, entity embedding size will be 2 * hidden_dim.
      Default: False
  double_relation_emb : bool
      If True, relation embedding size will be 2 * hidden_dim.
      Default: False
code: def __init__(
      self,
      args,
      model_name,
      n_entities,
      n_relations,
      hidden_dim,
      gamma,
      double_entity_emb=False,
      double_relation_emb=False,
      ent_feat_dim=-1,
      rel_feat_dim=-1,
  ):
      super(KEModel, self).__init__()
      self.args = args
      self.has_edge_importance = args.has_edge_importance
      self.n_entities = n_entities
      self.n_relations = n_relations
      self.model_name = model_name
      self.hidden_dim = hidden_dim
      self.eps = EMB_INIT_EPS
      self.emb_init = (gamma + self.eps) / hidden_dim
      entity_dim = 2 * hidden_dim if double_entity_emb else hidden_dim
      relation_dim = (
          2 * hidden_dim if double_relation_emb else hidden_dim
      )
      self.encoder_model_name = args.encoder_model_name
  
      device = get_device(args)
  
      self.loss_gen = LossGenerator(
          args,
          args.loss_genre,
          args.neg_adversarial_sampling,
          args.adversarial_temperature,
          args.pairwise,
      )
  
      if self.encoder_model_name in ["shallow", "concat"]:
          self.entity_emb = ExternalEmbedding(
              args,
              n_entities,
              entity_dim,
              F.cpu() if args.mix_cpu_gpu else device,
          )
      if self.encoder_model_name in ["roberta", "concat"]:
          assert ent_feat_dim != -1 and rel_feat_dim != -1
          self.entity_feat = ExternalEmbedding(
              args,
              n_entities,
              ent_feat_dim,
              F.cpu() if args.mix_cpu_gpu else device,
              is_feat=True,
          )
      # For RESCAL, relation_emb = relation_dim * entity_dim
      if model_name == "RESCAL":
          rel_dim = relation_dim * entity_dim
      else:
          rel_dim = relation_dim
  
      self.use_mlp = self.encoder_model_name in ["concat", "roberta"]
      if self.encoder_model_name == "concat":
          self.transform_net = MLP(
              entity_dim + ent_feat_dim,
              entity_dim,
              relation_dim + rel_feat_dim,
              relation_dim,
          )
          # self.transform_e_net = torch.nn.Linear(entity_dim, entity_dim)
          # self.transform_r_net = torch.nn.Linear(relation_dim, relation_dim)
      elif self.encoder_model_name == "roberta":
          self.transform_net = MLP(
              ent_feat_dim, entity_dim, rel_feat_dim, relation_dim
          )
  
      self.rel_dim = rel_dim
      self.entity_dim = entity_dim
      self.strict_rel_part = args.strict_rel_part
      self.soft_rel_part = args.soft_rel_part
      print(self.strict_rel_part, self.soft_rel_part)
      assert not self.strict_rel_part and not self.soft_rel_part
      if not self.strict_rel_part and not self.soft_rel_part:
          if self.encoder_model_name in ["shallow", "concat"]:
              self.relation_emb = ExternalEmbedding(
                  args,
                  n_relations,
                  rel_dim,
                  F.cpu() if args.mix_cpu_gpu else device,
              )
          if self.encoder_model_name in ["roberta", "concat"]:
              self.relation_feat = ExternalEmbedding(
                  args,
                  n_relations,
                  rel_feat_dim,
                  F.cpu() if args.mix_cpu_gpu else device,
                  is_feat=True,
              )
      else:
          self.global_relation_emb = ExternalEmbedding(
              args, n_relations, rel_dim, F.cpu()
          )
  
      if model_name == "TransE" or model_name == "TransE_l2":
          self.score_func = TransEScore(gamma, "l2")
      elif model_name == "TransE_l1":
          self.score_func = TransEScore(gamma, "l1")
      elif model_name == "TransR":
          projection_emb = ExternalEmbedding(
              args,
              n_relations,
              entity_dim * relation_dim,
              F.cpu() if args.mix_cpu_gpu else device,
          )
  
          self.score_func = TransRScore(
              gamma, projection_emb, relation_dim, entity_dim
          )
      elif model_name == "DistMult":
          self.score_func = DistMultScore()
      elif model_name == "ComplEx":
          self.score_func = ComplExScore()
      elif model_name == "RESCAL":
          self.score_func = RESCALScore(relation_dim, entity_dim)
      elif model_name == "RotatE":
          self.score_func = RotatEScore(gamma, self.emb_init)
      elif model_name == "SimplE":
          self.score_func = SimplEScore()
  
      self.model_name = model_name
      self.head_neg_score = self.score_func.create_neg(True)
      self.tail_neg_score = self.score_func.create_neg(False)
      self.head_neg_prepare = self.score_func.create_neg_prepare(True)
      self.tail_neg_prepare = self.score_func.create_neg_prepare(False)
  
      self.reset_parameters()
  
imports: import os
  import numpy
  import math
  import dgl.backend
  import pdb
  import torch
  import torch.nn
  from .pytorch.tensor_models import logsigmoid
  from .pytorch.tensor_models import abs
  from .pytorch.tensor_models import masked_select
  from .pytorch.tensor_models import get_device, get_dev
  from .pytorch.tensor_models import norm
  from .pytorch.tensor_models import get_scalar
  from .pytorch.tensor_models import reshape
  from .pytorch.tensor_models import cuda
  from .pytorch.tensor_models import ExternalEmbedding
  from .pytorch.tensor_models import InferEmbedding
  from .pytorch.score_fun import *
  from .pytorch.loss import LossGenerator
label_desc: Use image node_id 1 to create a new KEModel object from inherited base classes: object with example: obj = KEModel(args, model_name, n_entities, n_relations, hidden_dim, gamma, double_entity_emb, double_relation_emb, ent_feat_dim, rel_feat_dim)