def repeat_interleave(input, repeats, dim=0):
    """
    Repeat and interleave a tensor along a given dimension. The current PyTorch
    implementation of `repeat_interleave` is slow and there is an open ticket for it:
    https://github.com/pytorch/pytorch/issues/31980
    Args:
        input (torch.Tensor): Tensor containing the input data.
        repeats (int): The number of repetitions for each element.
        dim (int). The dimension along which to repeat values.

    Returns:
        torch.Tensor: Repeated tensor which has the same shape as input, except along
            the given axis.
    """
    output = input.unsqueeze(dim + 1).expand(-1, repeats, *input.shape[1:])
    if dim == 0:
        output_shape = (-1,) + input.shape[1:]
    else:
        output_shape = input.shape[:dim] + (-1,) + input.shape[dim:]
    return output.reshape(output_shape)