def elementwise_or_scalar(*args):
    """Type propagation rule for operations like elementwise add/multiply,
    that allows arguments to be scalar. Rules:
    - All vectors in `args` must have the same shape and element type.
    - Any scalars in `args` must have the same type, and this type must
        match the element type of the vectors, if any.
    Example ShapeTypes and return ShapeType:
        ((), Float), ((2, 3), Tensor(2, Float))
            -> ((2, 3), Tensor(2, Float))
        ((2, 3), Tensor(2, Float)), ((2, 3), Tensor(2, Float))
            -> ((2, 3), Tensor(2, Float))
        ((2, 3), Tensor(2, Float)), ((), Int)
            -> ValueError
        ((2, 3), Tensor(2, Float)), ((2, 3), Tensor(2, Int)))
            -> ValueError
        ((3, 2), Tensor(2, Float)), ((2, 3), Tensor(2, Float))
            -> ValueError
    """

    def _get_type(x):
        """Gets type if scalar, and element type if tensor."""
        if x.is_tensor:
            return _get_type(x.tensor_elem_type)

        if x.is_scalar:
            return x

        raise ValueError(f"Argument has unsupported type {x}")

    assert len(args) > 0
    # All types must be equal:
    types = set(_get_type(a.shape_type.type) for a in args)
    if len(types) > 1:
        raise ValueError(f"Arguments have unsupported types: {types}")
    # If there are any non-scalars, then their shapes must be equal:
    shapes = set(a.shape_type.shape for a in args if a.shape_type.shape != ())
    if len(shapes) > 1:
        raise ValueError(f"Arguments have incompatible shapes: {shapes}")
    # Return ShapeType is that of the vector, if there is one:
    try:
        return next(a.shape_type for a in args if a.shape_type.shape != ())
    except StopIteration:
        return args[0].shape_type