code: def latent_prediction_model(
      inputs,
      ed_attention_bias,
      latents_discrete,
      latents_dense,
      hparams,
      vocab_size=None,
      name=None,
  ):
      """Transformer-based latent prediction model.
  
      It is an autoregressive decoder over latents_discrete given inputs.
  
      Args:
        inputs: Tensor of shape [batch, length_kv, hparams.hidden_size]. Inputs to
          attend to for the decoder on latents.
        ed_attention_bias: Tensor which broadcasts with shape [batch,
          hparams.num_heads, length_q, length_kv]. Encoder-decoder attention bias.
        latents_discrete: Tensor of shape [batch, length_q, vocab_size].
          One-hot latents to compute log-probability of given inputs.
        latents_dense: Tensor of shape [batch, length_q, hparams.hidden_size].
          length_q is the latent length, which is
          height * width * hparams.num_latents / (2**hparams.num_compress_steps).
        hparams: HParams.
        vocab_size: int or None. If None, it is 2**hparams.bottleneck_bits.
        name: string, variable scope.
  
      Returns:
        latents_pred: Tensor of shape [batch, length_q, hparams.hidden_size].
        latents_pred_loss: Tensor of shape [batch, length_q].
      """
      with tf.variable_scope(name, default_name="latent_prediction"):
          if hparams.mode != tf_estimator.ModeKeys.PREDICT:
              latents_pred = transformer_latent_decoder(
                  tf.stop_gradient(latents_dense),
                  inputs,
                  ed_attention_bias,
                  hparams,
                  name,
              )
              if vocab_size is None:
                  vocab_size = 2**hparams.bottleneck_bits
              if not hparams.soft_em:
                  # TODO(trandustin): latents_discrete is not one-hot from
                  # discrete_bottleneck unless hparams.soft_em is True. Refactor.
                  latents_discrete = tf.one_hot(
                      latents_discrete, depth=vocab_size
                  )
              _, latent_pred_loss = ae_latent_softmax(
                  latents_pred,
                  tf.stop_gradient(latents_discrete),
                  vocab_size,
                  hparams,
              )
      return latents_pred, latent_pred_loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from six.moves import range
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import transformer_layers
  from tensor2tensor.utils import beam_search
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
label_desc: Use image node_id 17 for calling a global function with example usage: latent_prediction_model(inputs, ed_attention_bias, latents_discrete, latents_dense, hparams, vocab_size, name) and returns: latents_pred,  latent_pred_loss