def map_constraint_logits(self, original_ids,
                            constraint_ids,
                            constraint_logits):
    """Map constraint logits to original if ids match.

    constraint_logits is logits for constraint_ids. If an id in constrain_ids
    also appear in original_ids, we will map its logit to the corresponding
    position as in original_ids.

    Args:
      original_ids: order to map to
      constraint_ids: order to map from
      constraint_logits: logits in the same order as constraint_ids

    Returns:
      logits in the order of original_ids if exist, otherwise 0.
    """
    constraint_logits_mapping, _ = util.compute_x_in_set(
        constraint_ids, original_ids)
    # batch_size, topk * 2, topk
    constraint_logits_mapping = tf.cast(
        constraint_logits_mapping, dtype=tf.float32)
    # batch_size, topk * 2, topk
    mapped_constraint_logits = tf.matmul(
        constraint_logits_mapping, tf.expand_dims(constraint_logits, axis=2))
    # batch_size, topk * 2, 1
    mapped_constraint_logits = tf.squeeze(mapped_constraint_logits, axis=2)
    # batch_size, topk * 2
    return mapped_constraint_logits
Map constraint logits to original if ids match.

    constraint_logits is logits for constraint_ids. If an id in constrain_ids
    also appear in original_ids, we will map its logit to the corresponding
    position as in original_ids.

    Args:
      original_ids: order to map to
      constraint_ids: order to map from
      constraint_logits: logits in the same order as constraint_ids

    Returns:
      logits in the order of original_ids if exist, otherwise 0.
    
language/emql/model.py