code: def slice_hidden(x, hidden_size, num_blocks):
      """Slice encoder hidden state under num_blocks.
  
      Args:
        x: Encoder hidden state of shape [batch_size, latent_dim, hidden_size].
        hidden_size: Dimension of the latent space.
        num_blocks: Number of blocks in DVQ.
  
      Returns:
        Sliced states of shape [batch_size, latent_dim, num_blocks, block_dim].
      """
      batch_size, latent_dim, _ = common_layers.shape_list(x)
      block_dim = hidden_size // num_blocks
      x_sliced = tf.reshape(
          x, shape=[batch_size, latent_dim, num_blocks, block_dim]
      )
      return x_sliced
  
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 2 for calling a global function with example usage: slice_hidden(x, hidden_size, num_blocks) and returns: x_sliced