code: def ae_latent_sample_beam(
      latents_dense_in, inputs, ed, embed, hparams
  ):
      """Sample from the latent space in the autoencoder."""
  
      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,
                  )
              )
              latents_pred = decode_transformer(
                  inputs, ed, latents_dense, hparams, "extra"
              )
              logits = tf.layers.dense(
                  latents_pred,
                  2**hparams.bottleneck_bits,
                  name="extra_logits",
              )
              current_output_position = (
                  common_layers.shape_list(ids)[1] - 1
              )
              logits = logits[:, current_output_position, :, :]
          return tf.squeeze(logits, axis=[1])
  
      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,
          beam_size=1,
          decode_length=length,
          vocab_size=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
  import functools
  from six.moves import range
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import beam_search
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 11 for calling a global function with example usage: ae_latent_sample_beam(latents_dense_in, inputs, ed, embed, hparams) and returns: res, tf