def test_reduce(self): self.assertEqual([1, 1, 1, 1, 5], infer.reduce([1, 2, 3, 4, 5], [1, 2, 3])) self.assertEqual([1, 5], infer.reduce([1, 2, 3, 4, 5], [1, 2, 3], squeeze=True)) self.assertEqual([], infer.reduce([5], [-1], squeeze=True)) self.assertEqual([1, 6], infer.reduce([5, 6], [-2]))
def propagate_reduce( op, # type: ONNXOperation multi_axis, # type: bool default=None, # type: typing.Any dtype=None # type: typing.Optional[str] ): # type: (...)->typing.Tuple[typing.List[typing.List[int]], typing.List[str]] if default is None: default = list(range(op.input.rank)) if multi_axis: if 'axes' in op.attribs: axes = list(op.attribs['axes']) else: axes = list(default) else: if 'axis' in op.attribs: axes = [op.attribs['axis']] else: axes = list(default) return [ infer.reduce(input=op.input.shape, axes=axes, squeeze=not op.attribs.get('keepdims', 1)) ], [dtype if dtype is not None else op.input.dtype]
def propagate_reduce(op, const_value_by_tensor, dtype=None): # type: (TFOperation, _ConstValueByTensorT, str)->typing.Tuple[typing.List[typing.List[int]], typing.List[str]] input, axis = op.inputs axis = const_value_by_tensor[axis].tolist() # type: typing.List[int] if not isinstance(axis, list): axis = [axis] return [infer.reduce(input=input.shape, axes=axis, squeeze=not op.attribs["keep_dims"])], \ [op.attribs['T'] if not dtype else dtype]
def sum_reduce_like_shape(op): # type: (Caffe2Operation)->ShapeResult input, reference = op.inputs axis = op.attribs['axis'] target = [ None if i < axis or i >= axis + reference.rank else min( input.shape[i], reference.shape[i - axis]) for i in range(input.rank) ] axes_keep = [ i for i in range(input.rank) if target[i] == 1 and input.shape[i] != 1 ] axes_squeeze = [i for i in range(input.rank) if target[i] is None] tmp_shape = infer.reduce(input.shape, axes_keep, squeeze=False) shape = infer.reduce(tmp_shape, axes_squeeze, squeeze=True) return shape, input.dtype
def pool_shape(op): # type: (Caffe2Operation)->ShapeResult is_nhwc = op.attribs.get('order', 'NCHW').upper() == 'NHWC' if op.attribs['global_pooling']: return infer.reduce(op.inputs[0].shape, axes=list(range(1, op.inputs[0].rank - 1)) if is_nhwc else list(range(2, op.inputs[0].rank)), squeeze=False), op.inputs[0].dtype def expand(list, default): if is_nhwc: return [default] + list + [default] else: return [default, default] + list return infer.sliding_window( input=op.inputs[0].shape, filter=expand(op.attribs['kernels'], 1), padding=expand(caffe2_pads_to_nnef_padding(op.attribs['pads']), (0, 0)), stride=expand(op.attribs['strides'], 1), dilation=expand(op.attribs['dilations'], 1)), op.inputs[0].dtype
def propagate_argminmax(op, const_value_by_tensor): # type: (TFOperation, _ConstValueByTensorT)->typing.Tuple[typing.List[typing.List[int]], typing.List[str]] input, axis = op.inputs axis = const_value_by_tensor[axis].tolist() # type: int return [infer.reduce(input=input.shape, axes=[axis], squeeze=True)], [op.attribs['T']]
def propagate_reduce(op, const_value_by_tensor): # type: (TFOperation, _ConstValueByTensorT)->typing.Tuple[typing.List[typing.List[int]], typing.List[str]] input, axis = op.inputs axis = const_value_by_tensor[axis].tolist() # type: typing.List[int] return [infer.reduce(input=input.shape, axes=axis, squeeze=not op.attribs["keep_dims"])], [op.attribs['T']]
def reduce_shape(op, dtype=None): # type: (Caffe2Operation, typing.Optional[str])->ShapeResult axes = op.attribs['axes'] keep_dims = op.attribs['keepdims'] return infer.reduce(op.inputs[0].shape, axes=axes, squeeze=not keep_dims), \ op.inputs[0].dtype if dtype is None else dtype
def arg_min_max_shape(op, dtype=None): # type: (Caffe2Operation, typing.Optional[str])->ShapeResult axis = op.attribs.get('axis', -1) keep_dims = op.attribs.get('keepdims', 1) return infer.reduce(op.inputs[0].shape, axes=[axis], squeeze=not keep_dims), \ op.inputs[0].dtype if dtype is None else dtype
def reduction_shape(op): axis = op.attribs['axis'] rank = len(op.inputs[0].shape) axes = list(range(axis + rank if axis < 0 else axis, rank)) return shapes.reduce(op.inputs[0].shape, axes=axes, squeeze=True)