code: def embed(
      x,
      hidden_size,
      z_size,
      filter_size,
      bottleneck_kind="dvq",
      soft_em=False,
      num_blocks=2,
      num_residuals=1,
      block_v_size=None,
      means=None,
      name=None,
  ):
      """Embedding function that takes discrete latent and returns embedding.
  
      Args:
        x: Input to the discretization bottleneck.
        hidden_size: Dimension of the latent state.
        z_size: Number of bits, where discrete codes range from 1 to 2**z_size.
        filter_size: Dimension to project embedding by. Used only if bottleneck_kind
          is semhash.
        bottleneck_kind: Kind of discretization bottleneck to use; one of dvq,
          semhash, gumbel-softmax (Default: dvq).
        soft_em: If True then it uses a multi-sample version of EM (Default: False).
        num_blocks: Number of blocks in DVQ (Default: 2).
        num_residuals: Number of residuals (Default: 1).
        block_v_size: Number of embedding entries per block (Default: None).
        means: The embedding table for dvq (Default: None).
        name: Name for the bottleneck scope.
  
      Returns:
        Continuous embedding to be passed on to the decoder.
  
      Raises:
        ValueError: For unknown or missing arguments.
      """
      with tf.variable_scope(
          name, default_name="embed", reuse=tf.AUTO_REUSE
      ):
          if bottleneck_kind == "semhash":
              c = int_to_bit(x, z_size)
              h1a = tf.layers.dense(c, filter_size, name="vch1a")
              h1b = tf.layers.dense(1.0 - c, filter_size, name="vch1b")
              h1 = h1a + h1b
              h1 = tf.layers.dense(
                  h1, hidden_size, name="vch_final_linear"
              )
  
          elif bottleneck_kind == "gumbel-softmax":
              hot = tf.one_hot(x, 2**z_size)
              h1 = tf.layers.dense(hot, hidden_size, name="dae_dense")
          elif bottleneck_kind in ["dvq", "gumbel-softmax-dvq"]:
              if block_v_size is None:
                  raise ValueError(
                      "Bottleneck kind is dvq but block_v_size is None."
                  )
  
              if soft_em:
                  assert num_residuals == 1
                  x_hot_flat = tf.reshape(
                      x, shape=[-1, num_blocks, block_v_size]
                  )
                  h1 = tf.matmul(
                      tf.transpose(x_hot_flat, perm=[1, 0, 2]), means[0]
                  )
                  h1 = tf.transpose(h1, perm=[1, 0, 2])
                  new_shape = common_layers.shape_list(x)
                  new_shape[-1] = hidden_size
                  h1 = tf.reshape(h1, shape=new_shape)
              else:
                  shape_x = common_layers.shape_list(x)
                  x_flat = tf.reshape(x, [-1, 1])
                  c = int_to_bit(x_flat, num_bits=z_size, base=2)
                  shape = common_layers.shape_list(c)
                  new_shape = shape
                  new_shape[-1] = num_residuals
                  new_shape.append(num_blocks)
                  new_shape.append(
                      int(z_size / (num_residuals * num_blocks))
                  )
                  c = tf.to_int32(tf.reshape(c, shape=new_shape))
                  h1_shape = shape_x
                  h1_shape.append(hidden_size)
                  h1 = tf.zeros(dtype=tf.float32, shape=h1_shape)
                  for i in range(num_residuals):
                      c_residual = bit_to_int(
                          c[:, :, i, :, :],
                          num_bits=int(
                              z_size / (num_residuals * num_blocks)
                          ),
                          base=2,
                      )
                      c_hot = tf.one_hot(
                          c_residual, depth=block_v_size, axis=-1
                      )
                      c_hot_flat = tf.reshape(
                          c_hot, shape=[-1, num_blocks, block_v_size]
                      )
                      h1_residual = tf.matmul(
                          tf.transpose(c_hot_flat, perm=[1, 0, 2]),
                          means[i],
                      )
                      h1_residual = tf.transpose(
                          h1_residual, perm=[1, 0, 2]
                      )
                      h1_residual = tf.reshape(
                          h1_residual, shape=h1_shape
                      )
                      h1 += h1_residual
          elif bottleneck_kind == "rounding":
              h1 = x
          else:
              raise ValueError("Unknown bottleneck kind.")
  
          return h1
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from functools import partial
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 8 for calling a global function with example usage: embed(x, hidden_size, z_size, filter_size, bottleneck_kind, soft_em, num_blocks, num_residuals, block_v_size, means, name) and returns: h1