def compute_max_pool_embedding(input_embeddings, input_lengths):
  """Computes max pool embedding.

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

  Returns:
    max_pool_embedding: <tf.float32>[bs, emb_dim]
  """
  max_seq_len = tf.shape(input_embeddings)[1]
  # <tf.float32>[bs, max_seq_len]
  mask = 1.0 - tf.sequence_mask(input_lengths, max_seq_len, dtype=tf.float32)
  mask = tf.squeeze(mask * (-1e-6), 1)
  mask = tf.expand_dims(mask, 2)
  # <tf.float32>[bs, emb_dim]
  max_pool_embedding = tf.reduce_max(input_embeddings + mask, 1)
  # <tf.float32>[bs, dim]
  return max_pool_embedding
Computes max pool embedding.

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

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