def get_nonzero_ids(v_hot, k):
  """Get k ids of nonzero values in the v_hot vector.

  If less than k nonzero values exist, pad with -1.

  Args:
    v_hot: v hot vector (batch_size, dim)
    k: top k

  Returns:
    k ids of nonzero values
  """
  if k == -1:
    k = tf.shape(v_hot)[-1]
  label_k, idx_k = tf.nn.top_k(v_hot, k)
  idx_k = tf.where(
      tf.not_equal(label_k, 0), idx_k, tf.fill(tf.shape(idx_k), value=-1))
  return idx_k
Get k ids of nonzero values in the v_hot vector.

  If less than k nonzero values exist, pad with -1.

  Args:
    v_hot: v hot vector (batch_size, dim)
    k: top k

  Returns:
    k ids of nonzero values
  
language/emql/util.py