def embedding_lookup(embedding_table, input_ids, use_one_hot_embeddings=False):
  """Looks up embeddings for id tensor.

  Args:
    embedding_table: float Tensor of shape [vocab_size, embedding_size]
    input_ids: int32 Tensor of shape [batch_size, seq_length] or [batch_size,
      seq_length, attributes_size] containing ids.
    use_one_hot_embeddings: bool. If True, use one-hot method for word
      embeddings. If False, use `tf.nn.embedding_lookup()`.

  Returns:
    float Tensor of shape [batch_size, seq_length, embedding_size] or
    [batch_size, seq_length, attributes_size, embedding_size].
  """
  if not use_one_hot_embeddings:
    return tf.nn.embedding_lookup(embedding_table, input_ids)

  vocab_size, embedding_size = embedding_table.get_shape()
  flat_input_ids = tf.reshape(input_ids, [-1])
  one_hot_input_ids = tf.one_hot(flat_input_ids, depth=vocab_size)
  output = tf.matmul(one_hot_input_ids, embedding_table)

  input_shape = get_shape_list(input_ids)

  output = tf.reshape(output, input_shape + [embedding_size])
  return output
Looks up embeddings for id tensor.

  Args:
    embedding_table: float Tensor of shape [vocab_size, embedding_size]
    input_ids: int32 Tensor of shape [batch_size, seq_length] or [batch_size,
      seq_length, attributes_size] containing ids.
    use_one_hot_embeddings: bool. If True, use one-hot method for word
      embeddings. If False, use `tf.nn.embedding_lookup()`.

  Returns:
    float Tensor of shape [batch_size, seq_length, embedding_size] or
    [batch_size, seq_length, attributes_size, embedding_size].
  
language/xsp/model/common_layers.py