# Copyright (c) 2017-2019 Uber Technologies, Inc. # SPDX-License-Identifier: Apache-2.0 def broadcast_shape(*shapes, **kwargs): """ Similar to ``np.broadcast()`` but for shapes. Equivalent to ``np.broadcast(*map(np.empty, shapes)).shape``. :param tuple shapes: shapes of tensors. :param bool strict: whether to use extend-but-not-resize broadcasting. :returns: broadcasted shape :rtype: tuple :raises: ValueError """ strict = kwargs.pop("strict", False) reversed_shape = [] for shape in shapes: for i, size in enumerate(reversed(shape)): if i >= len(reversed_shape): reversed_shape.append(size) elif reversed_shape[i] == 1 and not strict: reversed_shape[i] = size elif reversed_shape[i] != size and (size != 1 or strict): raise ValueError( "shape mismatch: objects cannot be broadcast to a single shape: {}".format( " vs ".join(map(str, shapes)) ) ) return tuple(reversed(reversed_shape))