def project_hidden(x, projection_tensors, hidden_size, num_blocks):
  """Project encoder hidden state under num_blocks using projection tensors.

  Args:
    x: Encoder hidden state of shape [batch_size, latent_dim,  hidden_size].
    projection_tensors: Projection tensors used to project the hidden state.
    hidden_size: Dimension of the latent space.
    num_blocks: Number of blocks in DVQ.

  Returns:
    x_projected: Projected states of shape [batch_size, latent_dim, num_blocks,
      hidden_size / num_blocks].
  """
  batch_size, latent_dim, _ = common_layers.shape_list(x)
  x = tf.reshape(x, shape=[1, -1, hidden_size])
  x_tiled = tf.reshape(
      tf.tile(x, multiples=[num_blocks, 1, 1]),
      shape=[num_blocks, -1, hidden_size])
  x_projected = tf.matmul(x_tiled, projection_tensors)
  x_projected = tf.transpose(x_projected, perm=[1, 0, 2])
  x_4d = tf.reshape(x_projected, [batch_size, latent_dim, num_blocks, -1])
  return x_4d
def project_hidden ( x , projection_tensors , hidden_size , num_blocks ) : batch_size , latent_dim , _ = common_layers . shape_list ( x ) x = tf . reshape ( x , shape = [ 1 , - 1 , hidden_size ] ) x_tiled = tf . reshape ( tf . tile ( x , multiples = [ num_blocks , 1 , 1 ] ) , shape = [ num_blocks , - 1 , hidden_size ] ) x_projected = tf . matmul ( x_tiled , projection_tensors ) x_projected = tf . transpose ( x_projected , perm = [ 1 , 0 , 2 ] ) x_4d = tf . reshape ( x_projected , [ batch_size , latent_dim , num_blocks , - 1 ] ) return x_4d
Project encoder hidden state under num_blocks using projection tensors.

  Args:
    x: Encoder hidden state of shape [batch_size, latent_dim,  hidden_size].
    projection_tensors: Projection tensors used to project the hidden state.
    hidden_size: Dimension of the latent space.
    num_blocks: Number of blocks in DVQ.

  Returns:
    x_projected: Projected states of shape [batch_size, latent_dim, num_blocks,
      hidden_size / num_blocks].