def compute_x_in_set(x, s):
  """Check if elements in tensor x are in set s.

  Args:
    x: batch_size, num_candidate
    s: batch_size, k (padded with -1)

  Returns:
    boolean tensor with shape of x
  """
  s = tf.expand_dims(s, axis=2)
  # batch_size, k, 1
  k = tf.shape(s)[1]
  x = tf.tile(tf.expand_dims(x, axis=1), [1, k, 1])
  # batch_size, k, num_candidate
  x_to_s_map = tf.equal(x, s)
  # batch_size, k, num_candidate
  x_in_s = tf.reduce_any(x_to_s_map, axis=1)
  # batch_size, num_candidate
  return x_to_s_map, x_in_s
Check if elements in tensor x are in set s.

  Args:
    x: batch_size, num_candidate
    s: batch_size, k (padded with -1)

  Returns:
    boolean tensor with shape of x
  
language/emql/util.py