code: 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
  
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 1 for calling a global function with example usage: project_hidden(x, projection_tensors, hidden_size, num_blocks) and returns: x_4d