import torch


def split_stack(x, split_sizes, split_dim, stack_dim):
    """Split x along dimension split_dim and stack again at dimension stack_dim"""
    t = torch.stack(torch.split(x, split_sizes, dim=split_dim), dim=stack_dim)
    return t