code: def ae_latent_softmax(latents_pred, latents_discrete, hparams):
      """Latent prediction and loss."""
      vocab_size = 2**hparams.z_size
      if hparams.num_decode_blocks < 2:
          latents_logits = tf.layers.dense(
              latents_pred, vocab_size, name="extra_logits"
          )
          if hparams.logit_normalization:
              latents_logits *= tf.rsqrt(
                  1e-8 + tf.reduce_mean(tf.square(latents_logits))
              )
  
          loss = None
          if latents_discrete is not None:
              if hparams.soft_em:
                  # latents_discrete is actually one-hot of multinomial samples
                  assert hparams.num_decode_blocks == 1
                  loss = tf.nn.softmax_cross_entropy_with_logits_v2(
                      labels=latents_discrete, logits=latents_logits
                  )
              else:
                  loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
                      labels=latents_discrete, logits=latents_logits
                  )
          sample = multinomial_sample(
              latents_logits, vocab_size, hparams.sampling_temp
          )
          return sample, loss
  
      # Multi-block case.
      vocab_bits = int(math.log(vocab_size, 2))
      assert vocab_size == 2**vocab_bits
      assert vocab_bits % hparams.num_decode_blocks == 0
      block_vocab_size = 2 ** (vocab_bits // hparams.num_decode_blocks)
      latents_logits = [
          tf.layers.dense(
              latents_pred, block_vocab_size, name="extra_logits_%d" % i
          )
          for i in range(hparams.num_decode_blocks)
      ]
      loss = None
      if latents_discrete is not None:
          losses = []
          for i in range(hparams.num_decode_blocks):
              d = tf.floormod(
                  tf.floordiv(latents_discrete, block_vocab_size**i),
                  block_vocab_size,
              )
              losses.append(
                  tf.nn.sparse_softmax_cross_entropy_with_logits(
                      labels=d, logits=latents_logits[i]
                  )
              )
          loss = sum(losses)
      samples = [
          multinomial_sample(l, block_vocab_size, hparams.sampling_temp)
          for l in latents_logits
      ]
      sample = sum(
          [s * block_vocab_size**i for i, s in enumerate(samples)]
      )
      return sample, loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import math
  import os
  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 discretization
  from tensor2tensor.layers import latent_layers
  from tensor2tensor.layers import modalities
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import beam_search
  from tensor2tensor.utils import contrib
  from tensor2tensor.utils import expert_utils
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 10 for calling a global function with example usage: ae_latent_softmax(latents_pred, latents_discrete, hparams) and returns: sample,  loss, sample,  loss