class_docstr: Calculate simularity of entity/relation embeddings based on pretained KG embeddings
  Parameters
  ---------
  device : int
      Device to run the inference, -1 for CPU
  emb_file : dict
      Containing embedding information
  sfunc : str
      What kind of score is used,
          cosine: score = $rac{x \cdot y}{||x||_2||y||_2}$
          l2: score = $-||x - y||_2$
          l1: score = $-||x - y||_1$
          dot: score = $x \cdot y$
          ext_jaccard: score = $rac{x \cdot y}{||x||_{2}^{2} + ||y||_{2}^{2} - x \cdot y}$
class_docstr_tok: ["Calculate","simularity","of","entity\/relation","embeddings","based","on","pretained","KG","embeddings","Parameters","--","--","--","--","-","device",":","int","Device","to","run","the","inference",",","-1","for","CPU","emb_file",":","dict","Containing","embedding","information","sfunc",":","str","What","kind","of","score","is","used",",","cosine",":","score","=","$","rac","{","x","\\cdot","y","}","{","||x||_2||y||_2","}","$","l2",":","score","=","$","-||x","-","y||_2","$","l1",":","score","=","$","-||x","-","y||_1","$","dot",":","score","=","$","x","\\cdot","y","$","ext_jaccard",":","score","=","$","rac","{","x","\\cdot","y","}","{","||x||_","{","2","}","^","{","2","}","+","||y||_","{","2","}","^","{","2","}","-","x","\\cdot","y","}","$"]
code: def topK(
      self, head=None, tail=None, bcast=False, pair_ws=False, k=10
  ):
      if head is None:
          head = F.arange(0, self.emb.shape[0])
      else:
          head = F.tensor(head)
      if tail is None:
          tail = F.arange(0, self.emb.shape[0])
      else:
          tail = F.tensor(tail)
  
      head_emb = self.emb[head]
      tail_emb = self.emb[tail]
      if pair_ws is True:
          result = []
          batch_size = self.batch_size
          # chunked cal score
          score = []
          num_head = head.shape[0]
          num_tail = tail.shape[0]
          for i in range((num_head + batch_size - 1) // batch_size):
              sh_emb = head_emb[
                  i * batch_size : (i + 1) * batch_size
                  if (i + 1) * batch_size < num_head
                  else num_head
              ]
              sh_emb = F.copy_to(sh_emb, self.device)
              st_emb = tail_emb[
                  i * batch_size : (i + 1) * batch_size
                  if (i + 1) * batch_size < num_head
                  else num_head
              ]
              st_emb = F.copy_to(st_emb, self.device)
              score.append(
                  F.copy_to(
                      self.sim_func(sh_emb, st_emb, pw=True), F.cpu()
                  )
              )
          score = F.cat(score, dim=0)
  
          sidx = F.argsort(score, dim=0, descending=True)
          sidx = sidx[:k]
          score = score[sidx]
          result.append(
              (
                  F.asnumpy(head[sidx]),
                  F.asnumpy(tail[sidx]),
                  F.asnumpy(score),
              )
          )
      else:
          num_head = head.shape[0]
          num_tail = tail.shape[0]
          batch_size = self.batch_size
  
          # chunked cal score
          score = []
          for i in range((num_head + batch_size - 1) // batch_size):
              sh_emb = head_emb[
                  i * batch_size : (i + 1) * batch_size
                  if (i + 1) * batch_size < num_head
                  else num_head
              ]
              sh_emb = F.copy_to(sh_emb, self.device)
              s_score = []
              for j in range((num_tail + batch_size - 1) // batch_size):
                  st_emb = tail_emb[
                      j * batch_size : (j + 1) * batch_size
                      if (j + 1) * batch_size < num_tail
                      else num_tail
                  ]
                  st_emb = F.copy_to(st_emb, self.device)
                  s_score.append(
                      F.copy_to(self.sim_func(sh_emb, st_emb), F.cpu())
                  )
              score.append(F.cat(s_score, dim=1))
          score = F.cat(score, dim=0)
  
          if bcast is False:
              result = []
              idx = F.arange(0, num_head * num_tail)
              score = F.reshape(score, (num_head * num_tail,))
  
              sidx = F.argsort(score, dim=0, descending=True)
              sidx = sidx[:k]
              score = score[sidx]
              sidx = sidx
              idx = idx[sidx]
              tail_idx = idx % num_tail
              idx = floor_divide(idx, num_tail)
              head_idx = idx % num_head
  
              result.append(
                  (
                      F.asnumpy(head[head_idx]),
                      F.asnumpy(tail[tail_idx]),
                      F.asnumpy(score),
                  )
              )
  
          else:  # bcast at head
              result = []
              for i in range(num_head):
                  i_score = score[i]
  
                  sidx = F.argsort(i_score, dim=0, descending=True)
                  idx = F.arange(0, num_tail)
                  i_idx = sidx[:k]
                  i_score = i_score[i_idx]
                  idx = idx[i_idx]
  
                  result.append(
                      (
                          np.full((k,), F.asnumpy(head[i])),
                          F.asnumpy(tail[idx]),
                          F.asnumpy(i_score),
                      )
                  )
  
      return result
  
file_path: ogb/examples/lsc/wikikg90m/dgl-ke-ogb-lsc/python/dglke/models/infer.py
imports: import os
  import time
  import argparse
  import numpy
  import dgl.backend
  from .general_models import InferModel
label_desc: Use image node_id 3 for calling the EmbSimInfer obj's underlying member method code with example usage: obj.topK(head, tail, bcast, pair_ws, k) and returns: result