def _transform_output_shape(
    output: Tensor, shape: Tuple[int, ...], *, reference_shape: Optional[Tensor] = None
) -> Tensor:
    r"""Collapse the broadcasted batch dimensions an input tensor to be the specified shape.

    Args:
        output: Tensor
        shape: List/tuple of int
        reference_shape: Tensor representation of shape to control which dimensions are collapsed.

    Returns:
        Tensor

    """
    out_tensor = output.clone()

    for dim in range(len(out_tensor.shape) - len(shape)):
        idx = 0
        if reference_shape is not None and out_tensor.shape[0] == reference_shape[0] != 1 and len(shape) > 2:
            idx = 1
        if out_tensor.shape[idx] != 1:
            raise AssertionError(f"Dimension {dim} of input is expected to be 1, got {out_tensor.shape[idx]}")
        out_tensor = out_tensor.squeeze(idx)

    return out_tensor
Collapse the broadcasted batch dimensions an input tensor to be the specified shape.

    Args:
        output: Tensor
        shape: List/tuple of int
        reference_shape: Tensor representation of shape to control which dimensions are collapsed.

    Returns:
        Tensor

    
kornia/augmentation/utils/helpers.py