def model_query2box(self, name, features,
                      params):
    """Model query2box for 9 different tasks.

    This function simulate 9 compositional queries as defined in the
    Query2Box paper:
      1c: X.follow(R)
      2c: X.follow(R1).follow(R2)
      3c: X.follow(R1).follow(R2).follow(R3)
      2i: X1.follow(R1) & X2.follow(R2)
      3i: X1.follow(R1) & X2.follow(R2) & X3.follow(R3)
      ic: (X1.follow(R1) & X2.follow(R2)).follow(R3)
      ci: X1.follow(R1) & X2.follow(R2).follow(R3)
      2u: X1.follow(R1) | X2.follow(R2)
      uc: (X1.follow(R1) | X2.follow(R2)).follow(R3)
    Modules (follow, intersection, union, decode) are defined in module.py.

    Args:
      name: name of the task, e.g. "query2box_2i".
      features: a tuple of tf.Tensors. The order is decided by each task.
      params: a dictionary of hyper-parameters

    Returns:
      loss and predictions
    """
    task = name.split('_')[-1]

    if task == '1c':  # X.follow(R)
      ent, rel1 = features
      ent_embs, ent_sketch = self.simple_entity_encode_op(params, ent)
      rel1_embs, rel1_sketch = self.simple_relation_encode_op(params, rel1)

      answer_embs, answer_sketch = self.simple_follow_op(
          params, ent_embs, rel1_embs, ent_sketch, rel1_sketch)

    elif task == '2c':  # X.follow(R1).follow(R2)
      ent, rel1, rel2 = features
      ent_embs, ent_sketch = self.simple_entity_encode_op(params, ent)
      rel1_embs, rel1_sketch = self.simple_relation_encode_op(params, rel1)
      rel2_embs, rel2_sketch = self.simple_relation_encode_op(params, rel2)

      obj1_embs, obj1_sketch = self.simple_follow_op(
          params, ent_embs, rel1_embs, ent_sketch, rel1_sketch)
      answer_embs, answer_sketch = self.simple_follow_op(
          params, obj1_embs, rel2_embs, obj1_sketch, rel2_sketch)

    elif task == '3c':  # X.follow(R1).follow(R2).follow(R3)
      ent, rel1, rel2, rel3 = features
      ent_embs, ent_sketch = self.simple_entity_encode_op(params, ent)
      rel1_embs, rel1_sketch = self.simple_relation_encode_op(params, rel1)
      rel2_embs, rel2_sketch = self.simple_relation_encode_op(params, rel2)
      rel3_embs, rel3_sketch = self.simple_relation_encode_op(params, rel3)

      obj1_embs, obj1_sketch = self.simple_follow_op(
          params, ent_embs, rel1_embs, ent_sketch, rel1_sketch)
      obj2_embs, obj2_sketch = self.simple_follow_op(
          params, obj1_embs, rel2_embs, obj1_sketch, rel2_sketch)
      answer_embs, answer_sketch = self.simple_follow_op(
          params, obj2_embs, rel3_embs, obj2_sketch, rel3_sketch)

    elif task == '2i' or task == '2u':  # X1.follow(R1) & / | X2.follow(R2)
      ent1, rel1, ent2, rel2 = features
      ent1_embs, ent1_sketch = self.simple_entity_encode_op(params, ent1)
      rel1_embs, rel1_sketch = self.simple_relation_encode_op(params, rel1)
      ent2_embs, ent2_sketch = self.simple_entity_encode_op(params, ent2)
      rel2_embs, rel2_sketch = self.simple_relation_encode_op(params, rel2)

      obj1_embs, obj1_sketch = self.simple_follow_op(
          params, ent1_embs, rel1_embs, ent1_sketch, rel1_sketch)
      obj2_embs, obj2_sketch = self.simple_follow_op(
          params, ent2_embs, rel2_embs, ent2_sketch, rel2_sketch)
      if task == '2i':
        answer_embs, answer_sketch = module.intersection_op(
            obj1_embs, obj1_sketch, obj2_embs, obj2_sketch)
      elif task == '2u':
        answer_embs, answer_sketch = module.union_op(
            obj1_embs, obj1_sketch, obj2_embs, obj2_sketch)

    elif task == '3i':  # X1.follow(R1) & X2.follow(R2) & X3.follow(R3)
      ent1, rel1, ent2, rel2, ent3, rel3 = features
      ent1_embs, ent1_sketch = self.simple_entity_encode_op(params, ent1)
      rel1_embs, rel1_sketch = self.simple_relation_encode_op(params, rel1)
      ent2_embs, ent2_sketch = self.simple_entity_encode_op(params, ent2)
      rel2_embs, rel2_sketch = self.simple_relation_encode_op(params, rel2)
      ent3_embs, ent3_sketch = self.simple_entity_encode_op(params, ent3)
      rel3_embs, rel3_sketch = self.simple_relation_encode_op(params, rel3)

      obj1_embs, obj1_sketch = self.simple_follow_op(
          params, ent1_embs, rel1_embs, ent1_sketch, rel1_sketch)
      obj2_embs, obj2_sketch = self.simple_follow_op(
          params, ent2_embs, rel2_embs, ent2_sketch, rel2_sketch)
      obj3_embs, obj3_sketch = self.simple_follow_op(
          params, ent3_embs, rel3_embs, ent3_sketch, rel3_sketch)
      answer_embs, answer_sketch = module.intersection_op(
          obj1_embs, obj1_sketch, obj2_embs, obj2_sketch)
      answer_embs, answer_sketch = module.intersection_op(
          answer_embs, answer_sketch, obj3_embs, obj3_sketch)

    elif task == 'ic' or task == 'uc':
      # (X1.follow(R1) & / | X2.follow(R2)).follow(R3)
      ent1, rel1, ent2, rel2, rel3 = features
      ent1_embs, ent1_sketch = self.simple_entity_encode_op(params, ent1)
      rel1_embs, rel1_sketch = self.simple_relation_encode_op(params, rel1)
      ent2_embs, ent2_sketch = self.simple_entity_encode_op(params, ent2)
      rel2_embs, rel2_sketch = self.simple_relation_encode_op(params, rel2)
      rel3_embs, rel3_sketch = self.simple_relation_encode_op(params, rel3)

      obj1_embs, obj1_sketch = self.simple_follow_op(
          params, ent1_embs, rel1_embs, ent1_sketch, rel1_sketch)
      obj2_embs, obj2_sketch = self.simple_follow_op(
          params, ent2_embs, rel2_embs, ent2_sketch, rel2_sketch)
      if task == 'ic':
        answer_embs, answer_sketch = module.intersection_op(
            obj1_embs, obj1_sketch, obj2_embs, obj2_sketch)
      elif task == 'uc':
        answer_embs, answer_sketch = module.union_op(
            obj1_embs, obj1_sketch, obj2_embs, obj2_sketch)
      answer_embs, answer_sketch = self.simple_follow_op(
          params, answer_embs, rel3_embs, answer_sketch, rel3_sketch)

    elif task == 'ci':  # X1.follow(R1) & X2.follow(R2).follow(R3)
      ent1, rel1, rel2, ent2, rel3 = features
      ent1_embs, ent1_sketch = self.simple_entity_encode_op(params, ent1)
      rel1_embs, rel1_sketch = self.simple_relation_encode_op(params, rel1)
      rel2_embs, rel2_sketch = self.simple_relation_encode_op(params, rel2)
      ent2_embs, ent2_sketch = self.simple_entity_encode_op(params, ent2)
      rel3_embs, rel3_sketch = self.simple_relation_encode_op(params, rel3)

      obj1_embs, obj1_sketch = self.simple_follow_op(
          params, ent1_embs, rel1_embs, ent1_sketch, rel1_sketch)
      obj2_embs, obj2_sketch = self.simple_follow_op(
          params, obj1_embs, rel2_embs, obj1_sketch, rel2_sketch)
      obj3_embs, obj3_sketch = self.simple_follow_op(
          params, ent2_embs, rel3_embs, ent2_sketch, rel3_sketch)
      answer_embs, answer_sketch = module.intersection_op(
          obj2_embs, obj2_sketch, obj3_embs, obj3_sketch)

    else:
      raise ValueError('task: %s not recognized' % task)

    # Decode from set representation to a list of entities. We will apply a
    # null sketch for decoding at query time.
    answer_sketch = None
    answer_ids, unused_answer_logits = module.decode_op(
        params, answer_embs, answer_sketch,
        self.entity_embeddings_mat, self.all_entity_sketches)

    return tf.constant(0.0), {'answer_ids': answer_ids}
Model query2box for 9 different tasks.

    This function simulate 9 compositional queries as defined in the
    Query2Box paper:
      1c: X.follow(R)
      2c: X.follow(R1).follow(R2)
      3c: X.follow(R1).follow(R2).follow(R3)
      2i: X1.follow(R1) & X2.follow(R2)
      3i: X1.follow(R1) & X2.follow(R2) & X3.follow(R3)
      ic: (X1.follow(R1) & X2.follow(R2)).follow(R3)
      ci: X1.follow(R1) & X2.follow(R2).follow(R3)
      2u: X1.follow(R1) | X2.follow(R2)
      uc: (X1.follow(R1) | X2.follow(R2)).follow(R3)
    Modules (follow, intersection, union, decode) are defined in module.py.

    Args:
      name: name of the task, e.g. "query2box_2i".
      features: a tuple of tf.Tensors. The order is decided by each task.
      params: a dictionary of hyper-parameters

    Returns:
      loss and predictions
    
language/emql/model.py