def _combine_heads(x):
  """Inverse of split_heads.

  Args:
    x: a Tensor with shape [batch, num_heads, length, emb_size / num_heads]

  Returns:
    a Tensor with shape [batch, length, emb_size]
  """
  x = tf.transpose(x, [0, 2, 1, 3])
  old_shape = x.get_shape().dims
  a, b = old_shape[-2:]
  new_shape = old_shape[:-2] + [a * b]
  ret = tf.reshape(x, tf.concat([tf.shape(x)[:-2], [-1]], 0))
  ret.set_shape(new_shape)
  return ret
Inverse of split_heads.

  Args:
    x: a Tensor with shape [batch, num_heads, length, emb_size / num_heads]

  Returns:
    a Tensor with shape [batch, length, emb_size]
  
language/xsp/model/common_layers.py