def decode_op(params,
              set_embs, set_sketch,
              embeddings_mat, all_sketches
              ):
  """Decode set representations to member ids.

  Args:
    params: a dictionary of hyper-parameters.
    set_embs: dense embeddings of set.
    set_sketch: sparse count-min sketch of set.
    embeddings_mat: embedding lookup table for elements in the set.
    all_sketches: precomputed hash ids of all candidate elements.

  Returns:
    Ids of the top k predicted elements.
  """
  topk_logits, topk_ids = util.topk_search_fn(
      set_embs, embeddings_mat, params['intermediate_top_k'])
  # batch_size, top_k     batch_size, top_k
  if set_sketch is not None:
    is_topk_eligible, _ = check_topk_fact_eligible(
        topk_ids, set_sketch, all_sketches, params)
    topk_logits += (1 - is_topk_eligible) * _VERY_NEG
    # batch_size, top_k
  _, argmax = tf.nn.top_k(topk_logits, k=params['intermediate_top_k'])
  # batch_size, top_k
  topk_ids = tf.gather(topk_ids, argmax, batch_dims=1)
  # batch_size, top_k
  return topk_ids, topk_logits
Decode set representations to member ids.

  Args:
    params: a dictionary of hyper-parameters.
    set_embs: dense embeddings of set.
    set_sketch: sparse count-min sketch of set.
    embeddings_mat: embedding lookup table for elements in the set.
    all_sketches: precomputed hash ids of all candidate elements.

  Returns:
    Ids of the top k predicted elements.
  
language/emql/module.py