code: def ae_latent_softmax(latents_pred, latents_discrete_hot, vocab_size, hparams):
    
    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
docs: 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.