def compute_average_embedding(input_embeddings, input_lengths):
  """Computes bag-of-words embedding.

  Args:
    input_embeddings: <tf.float32>[bs, max_seq_len, emb_dim]
    input_lengths: <tf.int64>[bs, 1]

  Returns:
    bow_embedding: <tf.float32>[bs, emb_dim]
  """
  max_seq_len = tf.shape(input_embeddings)[1]
  # <tf.float32>[bs, 1, max_seq_len]
  mask = tf.sequence_mask(input_lengths, max_seq_len, dtype=tf.float32)
  # <tf.float32>[bs, 1, emb_dim]
  sum_embedding = tf.matmul(mask, input_embeddings)
  # <tf.float32>[bs, 1, emb_dim]
  avg_embedding = sum_embedding / tf.to_float(tf.expand_dims(input_lengths, 2))
  # <tf.float32>[bs, dim]
  return tf.squeeze(avg_embedding, 1)
Computes bag-of-words embedding.

  Args:
    input_embeddings: <tf.float32>[bs, max_seq_len, emb_dim]
    input_lengths: <tf.int64>[bs, 1]

  Returns:
    bow_embedding: <tf.float32>[bs, emb_dim]
  
tensor2tensor/models/neural_assistant.py