def lookup_key(query: tf.Tensor, key_vocab: tf.Tensor) -> tf.Tensor:
  """Look up the index of each element in query in key_vocab.

  Args:
    query: A `Tensor`.
    key_vocab: A 1-D `Tensor` of unique keys.

  Returns:
    The indices of the keys in query, determined by position in key_vocab.
  """

  def _lookup_key():
    # Obtain 0-indexed int64 positions for the keys in key_vocab.
    indices = tf.cast(tf.range(tf.size(key_vocab)), tf.int64)

    expanded_vocab_size = tf.expand_dims(tf.size(key_vocab), axis=0)
    matrix_shape = tf.concat([expanded_vocab_size, tf.shape(query)], axis=0)
    # Expand dims of key_vocab to rank of query.
    vocab_shape = tf.concat(
        [expanded_vocab_size,
         tf.ones(tf.rank(query), dtype=tf.int32)], axis=0)
    # Make copies of key_vocab to fill matrix_shape.
    expand_vocab = tf.broadcast_to(
        tf.reshape(key_vocab, vocab_shape), matrix_shape)
    # Make copies of indices to fill matrix_shape.
    expand_indices = tf.broadcast_to(
        tf.reshape(indices, vocab_shape), matrix_shape)
    # Make copies of query to fill matrix_shape.
    expand_query = tf.broadcast_to(query, matrix_shape)

    # Indices where expand_query equals expand_vocab is set to the key's
    # index. All the other indices are -1.
    expand_result = tf.where(
        tf.math.equal(expand_query, expand_vocab), expand_indices,
        tf.cast(tf.fill(matrix_shape, -1), tf.int64))
    # Reduce matrix above to desired 1-D shape.
    result = tf.math.reduce_max(expand_result, axis=0)
    result.set_shape(query.shape)
    return result

  def _check_vocab_size_and_lookup_key():
    return tf.cond(
        tf.math.equal(tf.size(key_vocab), 0),
        lambda: tf.cast(tf.fill(tf.shape(query), -1), tf.int64), _lookup_key)

  def _check_input_size_and_lookup_key():
    return tf.cond(
        tf.math.equal(tf.size(query),
                      0), lambda: tf.constant([], dtype=tf.int64),
        _check_vocab_size_and_lookup_key)

  return _check_input_size_and_lookup_key()
Look up the index of each element in query in key_vocab.

  Args:
    query: A `Tensor`.
    key_vocab: A 1-D `Tensor` of unique keys.

  Returns:
    The indices of the keys in query, determined by position in key_vocab.
  
tensorflow_transform/tf_utils.py