Function unified_dim_sizes

xarray/computation/apply_ufunc.py:621–642

(
    variables: Iterable[Variable], exclude_dims: AbstractSet = frozenset()
)

def unified_dim_sizes(
    variables: Iterable[Variable], exclude_dims: AbstractSet = frozenset()
) -> dict[Hashable, int]:
    dim_sizes: dict[Hashable, int] = {}

    for var in variables:
        if len(set(var.dims)) < len(var.dims):
            raise ValueError(
                "broadcasting cannot handle duplicate "
                f"dimensions on a variable: {list(var.dims)}"
            )
        for dim, size in zip(var.dims, var.shape, strict=True):
            if dim not in exclude_dims:
                if dim not in dim_sizes:
                    dim_sizes[dim] = size
                elif dim_sizes[dim] != size:
                    raise ValueError(
                        "operands cannot be broadcast together "
                        "with mismatched lengths for dimension "
                        f"{dim}: {dim_sizes[dim]} vs {size}"
                    )
    return dim_sizes


SLICE_NONE = slice(None)