def unflatten_blocks_nd(x, blocks_per_dimension):
  """Converts a flattened tensor into a normal blocked tensor.

  Args:
    x: a [batch, d1 * ... dn, items_in_block, depth] tensor
    blocks_per_dimension: a n-d list of integers for number of blocks in each
      dimension.

  Returns:
    a [batch, d1, d2, ..., dn, items_in_block, depth] tensor
  """
  x_shape = common_layers.shape_list(x)
  assert x_shape[1] == np.prod(blocks_per_dimension, dtype=np.int32)
  return tf.reshape(x, [-1] + list(blocks_per_dimension) + x_shape[-2:])
Converts a flattened tensor into a normal blocked tensor.

  Args:
    x: a [batch, d1 * ... dn, items_in_block, depth] tensor
    blocks_per_dimension: a n-d list of integers for number of blocks in each
      dimension.

  Returns:
    a [batch, d1, d2, ..., dn, items_in_block, depth] tensor
  
tensor2tensor/layers/common_attention.py