code: def ae_latent_sample_beam(
      latents_dense_in, inputs, ed, embed, hparams
  ):
      """Samples from the latent space in the autoencoder.
  
      Args:
        latents_dense_in: Tensor of shape [batch, length_q, ...]. Only the shape of
          its first two dimensions are used. length_q is the latent length, which is
          height * width * hparams.num_latents / (2**hparams.num_compress_steps).
        inputs: Tensor of shape [batch, length_kv, hparams.hidden_size]. Encodings
          to attend to in decoder.
        ed: Tensor which broadcasts with shape [batch, hparams.num_heads, length_q,
          length_kv]. Encoder-decoder attention bias.
        embed: Callable which embeds discrete latent hot-vectors and a hidden size
          and returns dense vectors.
        hparams: HParams.
  
      Returns:
        Tensor of shape [batch, length].
      """
  
      def symbols_to_logits_fn(ids):
          """Go from ids to logits."""
          ids = tf.expand_dims(
              ids, axis=2
          )  # Ids start with added all-zeros.
          latents_discrete = tf.pad(
              ids[:, 1:], [[0, 0], [0, 1], [0, 0]]
          )
  
          with tf.variable_scope(tf.get_variable_scope(), reuse=False):
              latents_dense = embed(
                  tf.one_hot(
                      latents_discrete,
                      depth=2**hparams.bottleneck_bits,
                  ),
                  hparams.hidden_size,
              )
              latents_pred = transformer_latent_decoder(
                  latents_dense,
                  inputs,
                  ed,
                  hparams,
                  name="latent_prediction",
              )
              logits = tf.layers.dense(
                  latents_pred,
                  2**hparams.bottleneck_bits,
                  name="logits_dense",
              )
              current_output_position = (
                  common_layers.shape_list(ids)[1] - 1
              )
              logits = logits[:, current_output_position, :]
          return logits
  
      initial_ids = tf.zeros(
          [tf.shape(latents_dense_in)[0]], dtype=tf.int32
      )
      length = tf.shape(latents_dense_in)[1]
      ids, _, _ = beam_search.beam_search(
          symbols_to_logits_fn,
          initial_ids,
          1,
          length,
          2**hparams.bottleneck_bits,
          alpha=0.0,
          eos_id=-1,
          stop_early=False,
      )
  
      res = tf.expand_dims(ids[:, 0, :], axis=2)  # Pick first beam.
      return res[:, 1:]
  
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 5 for calling a global function with example usage: ae_latent_sample_beam(latents_dense_in, inputs, ed, embed, hparams) and returns: res, logits