def load_kb_embeddings(self, name, params):
    """Construct or load KB embeddings.

    Args:
      name: task name
      params: params

    Returns:
      a tensor for kb embeddings
    """
    if (name in ['set_follow', 'metaqa2', 'metaqa3'] or
        name.startswith('query2box')):
      all_fact_subj_embs = tf.nn.embedding_lookup(self.entity_embeddings_mat,
                                                  self.all_fact_subjids)
      # num_facts, hidden_size
      all_fact_rel_embs = tf.nn.embedding_lookup(self.relation_embeddings_mat,
                                                 self.all_fact_relids)
      # num_facts, hidden_size
      all_fact_obj_embs = tf.nn.embedding_lookup(self.entity_embeddings_mat,
                                                 self.all_fact_objids)
      # num_facts, hidden_size
      kb_embs = tf.concat(
          [all_fact_subj_embs, all_fact_rel_embs, all_fact_obj_embs], axis=1)
      # num_fact, hidden_size * 3
      return kb_embs

    elif name in ['webqsp']:
      return util.load_db_checkpoint(
          var_name='embeddings_mat/kb_embeddings_mat',
          checkpoint_dir=params['kb_index'],
          cpu=False,
          trainable=False)
Construct or load KB embeddings.

    Args:
      name: task name
      params: params

    Returns:
      a tensor for kb embeddings
    
language/emql/model.py