code: def compute_last_embedding(input_embeddings, input_lengths, hparams):
      """Computes average of last K embedding.
  
      Args:
        input_embeddings: <tf.float32>[bs, max_seq_len, emb_dim]
        input_lengths: <tf.int64>[bs, 1]
        hparams: model hparams
  
      Returns:
        last_k_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
      )
      del_mask = tf.sequence_mask(
          input_lengths - hparams.last_k, max_seq_len, dtype=tf.float32
      )
      final_mask = mask - del_mask
      # <tf.float32>[bs, 1, emb_dim]
      sum_embedding = tf.matmul(final_mask, input_embeddings)
      # <tf.float32>[bs, 1, emb_dim]
      last_k_embedding = sum_embedding / tf.to_float(
          tf.expand_dims(
              tf.ones([tf.shape(input_embeddings)[0], 1])
              * hparams.last_k,
              2,
          )
      )
      # <tf.float32>[bs, dim]
      return tf.squeeze(last_k_embedding, 1)
  
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 1 for calling a global function with example usage: compute_last_embedding(input_embeddings, input_lengths, hparams) and returns: tf