def check_topk_fact_eligible(
    topk_subj_ids, query_entity_sketch,
    all_entity_sketches,
    params):
  """Extract probs from sketch.

  Args:
    topk_subj_ids: batch_size, topk
    query_entity_sketch: batch_size, depth, width
    all_entity_sketches: num_entities, depth
    params: hyper-parameters

  Returns:
    if topk_subj_ids are eligible and their probabilities
  """
  intermediate_top_k = tf.shape(topk_subj_ids)[1]
  topk_subj_ids = tf.reshape(topk_subj_ids, shape=[-1])
  # batch_size * topk
  topk_subj_sketches = tf.gather(all_entity_sketches, topk_subj_ids)
  # batch_size * topk, cm_depth
  topk_subj_sketches = tf.reshape(
      topk_subj_sketches,
      shape=[-1, intermediate_top_k, params['cm_depth']])
  # batch_size, topk, depth
  topk_fact_eligibility_bits = tf.gather(
      query_entity_sketch,
      tf.transpose(topk_subj_sketches, perm=[0, 2, 1]),
      batch_dims=2, axis=2)
  # batch_size, depth, topk
  topk_fact_prev_probs = tf.reduce_min(topk_fact_eligibility_bits, axis=1)
  # batch_size, topk
  is_topk_fact_eligible = tf.cast(
      topk_fact_prev_probs > 0.0, dtype=tf.float32)
  # batch_size, topk

  return is_topk_fact_eligible, topk_fact_prev_probs
Extract probs from sketch.

  Args:
    topk_subj_ids: batch_size, topk
    query_entity_sketch: batch_size, depth, width
    all_entity_sketches: num_entities, depth
    params: hyper-parameters

  Returns:
    if topk_subj_ids are eligible and their probabilities
  
language/emql/module.py