def split_heads(x, num_heads):
    """ Split heads

    :param x: A tensor with shape [batch, length, channels]
    :param num_heads: An integer

    :returns: A tensor with shape [batch, heads, length, channels / heads]
    """
    assert x.shape[-1] % num_heads == 0, str(x.shape)
    return x.reshape(x.shape[:-1] + (num_heads, x.shape[-1] // num_heads)).permute(0, 2, 1, 3)