import torch


def _extend_tensor(input_tensor: torch.Tensor, N: int) -> torch.Tensor:
    """
    Extend a tensor `input_tensor` with ndim > 2, `N` times along the batch
    dimension. This is done in the following sequence of steps (where `B` is
    the batch dimension):

    .. code-block:: python

        input_tensor (B, ...)
        -> add leading empty dimension (1, B, ...)
        -> expand (N, B, ...)
        -> reshape (N * B, ...)

    Args:
        input_tensor: torch.Tensor with ndim > 2 representing a batched input.
        N: number of times to extend each element of the batch.
    """
    if input_tensor.ndim < 2:
        raise ValueError("Input tensor must have ndimensions >= 2.")
    B = input_tensor.shape[0]
    non_batch_dims = tuple(input_tensor.shape[1:])
    constant_dims = (-1,) * input_tensor.ndim  # these dims are not expanded.
    return (
        input_tensor.clone()[None, ...]
        .expand(N, *constant_dims)
        .transpose(0, 1)
        .reshape(N * B, *non_batch_dims)
    )