def _split_heads(x, num_heads):
  """Split dimension 3 into multiple heads.

  Attempts to preserve static shape information.

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

  Returns:
    a Tensor with shape [batch, num_heads, length, emb_size / num_heads]
  """
  old_shape = x.get_shape().dims
  new_shape = old_shape[:-1] + [num_heads] + [old_shape[-1] // num_heads]
  ret = tf.reshape(x, tf.concat([tf.shape(x)[:-1], [num_heads, -1]], 0))
  ret.set_shape(new_shape)
  return tf.transpose(ret, [0, 2, 1, 3])
Split dimension 3 into multiple heads.

  Attempts to preserve static shape information.

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

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