def combine_heads_nd(x):
  """Inverse of split_heads_nd.

  Args:
    x: a [batch, num_heads, d1, ..., dn, depth // num_heads] tensor

  Returns:
    a [batch, d1, ...., dn, depth] tensor
  """
  num_dimensions = len(common_layers.shape_list(x)) - 3
  return combine_last_two_dimensions(
      tf.transpose(x, [0] + list(range(2, num_dimensions + 2)) +
                   [1, num_dimensions + 2]))
Inverse of split_heads_nd.

  Args:
    x: a [batch, num_heads, d1, ..., dn, depth // num_heads] tensor

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