code: 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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import six
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 2 for calling a global function with example usage: compute_max_pool_embedding(input_embeddings, input_lengths) and returns: max_pool_embedding