1# mypy: allow-untyped-defs 2import operator 3from functools import reduce 4 5 6def maybe_view(tensor, size, check_same_size=True): 7 if check_same_size and tensor.size() == size: 8 return tensor 9 return tensor.contiguous().view(size) 10 11 12def maybe_unexpand(tensor, old_size, check_same_size=True): 13 if check_same_size and tensor.size() == old_size: 14 return tensor 15 num_unsqueezed = tensor.dim() - len(old_size) 16 expanded_dims = [ 17 dim 18 for dim, (expanded, original) in enumerate( 19 zip(tensor.size()[num_unsqueezed:], old_size) 20 ) 21 if expanded != original 22 ] 23 24 for _ in range(num_unsqueezed): 25 tensor = tensor.sum(0, keepdim=False) 26 for dim in expanded_dims: 27 tensor = tensor.sum(dim, keepdim=True) 28 return tensor 29 30 31# Check whether the op enable broadcasting, and whether it is supported by ONNX. 32# If dims1 and dims2 are different, then broadcast is True. 33# We always assume the combination of dims1 and dims2 is broadcastable. 34# The following types of broadcasting are supported in ONNX: 35# 1) Only one element in dims2, such as dims2 = [1, 1] 36# 2) dims2 is suffix of dims1, such as dims1 = [2, 3, 4], and dims2 = [3, 4] 37# Details can be found here: https://github.com/onnx/onnx/blob/master/docs/Operators.md#Gemm 38def check_onnx_broadcast(dims1, dims2): 39 broadcast = False 40 supported = True 41 len1 = len(dims1) 42 len2 = len(dims2) 43 numel1 = reduce(operator.mul, dims1) 44 numel2 = reduce(operator.mul, dims2) 45 if len1 < len2: 46 broadcast = True 47 if numel2 != 1: 48 supported = False 49 elif len1 > len2: 50 broadcast = True 51 if numel2 != 1 and dims1[len1 - len2 :] != dims2: 52 supported = False 53 else: 54 if dims1 != dims2: 55 broadcast = True 56 if numel2 != 1: 57 supported = False 58 59 if not supported: 60 raise ValueError( 61 f"Numpy style broadcasting is not supported in ONNX. Input dims are: {dims1}, {dims2}" 62 ) 63 return broadcast 64