def _prepare_shape_for_expand_dims(shape, axes):
    """
    Creates the expanded new shape based on the shape and given axes

    Args:
        shape (tuple): the shape of the tensor
        axes Union(int, tuple(int), list(int)): the axes with dimensions expanded.

    Returns:
        new_shape(tuple): the shape with dimensions expanded.
    """

    new_shape = []
    shape_idx = 0
    new_shape_length = len(shape)

    # Convert to set
    if isinstance(axes, int):
        new_shape_length += 1
        if axes >= new_shape_length or axes < -new_shape_length:
            raise ValueError(
                f"axis {axes} is out of bounds for tensor of dimension {new_shape_length}")
        axes = {axes}

    elif isinstance(axes, (list, tuple)):
        new_shape_length += len(axes)
        for axis in axes:
            if axis >= new_shape_length or axis < -new_shape_length:
                raise ValueError(
                    f"axis {axis} is out of bounds for tensor of dimension {new_shape_length}")
        axes = set(axes)

    else:
        raise TypeError(
            f"only int, tuple and list are allowed for axes, but got {type(axes)}")

    for new_shape_idx in range(new_shape_length):
        if new_shape_idx in axes or new_shape_idx - new_shape_length in axes:
            new_shape.append(1)
        else:
            new_shape.append(shape[shape_idx])
            shape_idx += 1
    return tuple(new_shape)