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 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:]