import torch


def spatial_broadcast(slots, resolution):
    """
    :param slots: tensor of shape [batch_size, slots, features]
    :param resolution: tuple of integers (height, width)f
    :return: tensor of shape [batch_size*num_slots, height, width, features], the first dimension just reshaped
    and for height and width the array is copied multiple times
    """
    slots = slots.view(-1, slots.size(2))  # [(batch_size-2)*slots, features]
    slots = torch.unsqueeze(slots, 1)
    slots = torch.unsqueeze(slots, 1)  # [(batch_size-2)*slots,1, 1 features]
    slots = slots.expand(slots.size(0), resolution[0], resolution[1], slots.size(3))
    return slots