import torch


def merge_leading_dims(x, num_dims=2):
    """Reshapes the tensor `x` such that the first `num_dims` dimensions are merged to one."""
    new_shape = torch.Size([-1]) + x.shape[num_dims:]
    return torch.reshape(x, new_shape)