def pool(embed, seq_len):
  """Pool `embed` along the second dimension.

  Args:
    embed: <tf.float32>[batch, seq_len, dim]
    seq_len:  <tf.int32>[batch] sequence lengths of `embed`

  Returns:
    out: <tf.float32>[batch, dim]
  """
  attention_w = ops.last_dim_weighted_sum(embed, "w")
  attention_w = ops.mask_logits(attention_w, seq_len)
  attention_w = tf.expand_dims(tf.nn.softmax(attention_w), 1)

  # [batch, 1, len] * [batch, len, dim] -> [batch, 1, dim]
  return tf.squeeze(tf.matmul(attention_w, embed), 1)
Pool `embed` along the second dimension.

  Args:
    embed: <tf.float32>[batch, seq_len, dim]
    seq_len:  <tf.int32>[batch] sequence lengths of `embed`

  Returns:
    out: <tf.float32>[batch, dim]
  
language/boolq/run_recurrent_model_boolq.py