import torch


def tile_batch(x, multiplier, batch_dim=0):
    """Copy several times for all the tensors in a batch

    Given a batch of tensors whose shape is [..., bsz, ...], `tile_batch` copy all the tensors in this batch with
    k times (=> [..., bsz, k, ...]), and then restore its original dimension (=> [..., bsz*k, ...]).

    """
    x_size = x.size()
    out_size = x_size[:batch_dim] + (x_size[batch_dim] * multiplier,) + x_size[batch_dim + 1:]

    x_tiled = torch.unsqueeze(x, dim=batch_dim + 1)
    x_tiled = x_tiled.repeat(*[1 if d != batch_dim + 1 else multiplier for d in range(len(x_size) + 1)])
    x_tiled = x_tiled.view(*out_size)

    return x_tiled