code: def ae_latent_softmax(
      latents_pred, latents_discrete_hot, vocab_size, hparams
  ):
      """Latent prediction and loss.
  
      Args:
        latents_pred: Tensor of shape [..., depth].
        latents_discrete_hot: Tensor of shape [..., vocab_size].
        vocab_size: an int representing the vocab size.
        hparams: HParams.
  
      Returns:
        sample: Tensor of shape [...], a sample from a multinomial distribution.
        loss: Tensor of shape [...], the softmax cross-entropy.
      """
      with tf.variable_scope("latent_logits"):
          latents_logits = tf.layers.dense(
              latents_pred, vocab_size, name="logits_dense"
          )
          if hparams.logit_normalization:
              latents_logits *= tf.rsqrt(
                  1e-8 + tf.reduce_mean(tf.square(latents_logits))
              )
          loss = tf.nn.softmax_cross_entropy_with_logits_v2(
              labels=latents_discrete_hot, logits=latents_logits
          )
  
          # TODO(trandustin): tease this out from ae_latent_softmax.
          # we use just the loss portion to anchor prior / encoder on text.
          sample = multinomial_sample(
              latents_logits,
              vocab_size,
              hparams.sampling_method,
              hparams.sampling_temp,
          )
          return sample, 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 4 for calling a global function with example usage: ae_latent_softmax(latents_pred, latents_discrete_hot, vocab_size, hparams) and returns: sample,  loss