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
    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
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 ) : 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 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
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.