code: def ae_latent_sample(
      latents_dense, inputs, ed, embed, iters, hparams
  ):
      """Sample from the latent space in the autoencoder."""
      if hparams.num_decode_blocks < 2 and hparams.sampling_temp == 0.0:
          # TODO(lukaszkaiser): beam-search only works in non-blocked mode for now.
          tf.logging.info(
              "Running beam-search for latents with beam size 1."
          )
          return ae_latent_sample_beam(
              latents_dense, inputs, ed, embed, hparams
          )
      latents_pred = decode_transformer(
          inputs, ed, latents_dense, hparams, "extra"
      )
      latents_discrete, _ = ae_latent_softmax(
          latents_pred, None, hparams
      )
  
      def next_bit(latents_discrete, i):
          latents_discrete_prev = latents_discrete
          with tf.variable_scope(tf.get_variable_scope(), reuse=True):
              latents_dense = embed(latents_discrete)
              latents_pred = decode_transformer(
                  inputs, ed, latents_dense, hparams, "extra"
              )
              latents_discrete, _ = ae_latent_softmax(
                  latents_pred, None, hparams
              )
              return tf.concat(
                  [
                      latents_discrete_prev[:, : (i + 1), :],
                      latents_discrete[:, (i + 1) :, :],
                  ],
                  axis=1,
              )
  
      for i in range(iters):
          latents_discrete = next_bit(latents_discrete, i)
      return latents_discrete
  
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 12 for calling a global function with example usage: ae_latent_sample(latents_dense, inputs, ed, embed, iters, hparams) and returns: latents_discrete, ae_latent_sample_beam, tf