diff --git a/mindspore/python/mindspore/ops/_op_impl/tbe/__init__.py b/mindspore/python/mindspore/ops/_op_impl/tbe/__init__.py index 5b8fa48903b..08b852a9dfb 100644 --- a/mindspore/python/mindspore/ops/_op_impl/tbe/__init__.py +++ b/mindspore/python/mindspore/ops/_op_impl/tbe/__init__.py @@ -319,6 +319,7 @@ from .resize_nearest_neighbor_grad_ds import _resize_nearest_neighbor_grad_ds_tb from .pad_d import _pad_d_tbe from .pad_d_ds import _pad_d_ds_tbe from .arg_max_with_value import _arg_max_with_value_tbe +from .arg_max_with_value_ds import _arg_max_with_value_ds_tbe from .arg_min_with_value import _arg_min_with_value_tbe from .smooth_l1_loss import _smooth_l1_loss_tbe from .smooth_l1_loss_ds import _smooth_l1_loss_ds_tbe diff --git a/mindspore/python/mindspore/ops/_op_impl/tbe/arg_max_with_value_ds.py b/mindspore/python/mindspore/ops/_op_impl/tbe/arg_max_with_value_ds.py new file mode 100644 index 00000000000..87e1d9d2c27 --- /dev/null +++ b/mindspore/python/mindspore/ops/_op_impl/tbe/arg_max_with_value_ds.py @@ -0,0 +1,39 @@ +# Copyright 2022 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +"""ArgMaxWithValue op""" +from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType + +arg_max_with_value_ds_op_info = TBERegOp("ArgMaxWithValue") \ + .fusion_type("ELEMWISE") \ + .async_flag(False) \ + .binfile_name("arg_max_with_value.so") \ + .compute_cost(10) \ + .kernel_name("arg_max_with_value") \ + .partial_flag(True) \ + .dynamic_shape(True) \ + .attr("axis", "required", "int", "all") \ + .input(0, "x", False, "required", "all") \ + .output(0, "indice", False, "required", "all") \ + .output(1, "values", False, "required", "all") \ + .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.F16_Default) \ + .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.F32_Default) \ + .get_op_info() + + +@op_info_register(arg_max_with_value_ds_op_info) +def _arg_max_with_value_ds_tbe(): + """ArgMaxWithValue TBE register""" + return diff --git a/mindspore/python/mindspore/ops/operations/array_ops.py b/mindspore/python/mindspore/ops/operations/array_ops.py index d7a126d154e..613ee953dd5 100755 --- a/mindspore/python/mindspore/ops/operations/array_ops.py +++ b/mindspore/python/mindspore/ops/operations/array_ops.py @@ -355,6 +355,20 @@ class Cast(PrimitiveWithInfer): if 'min_shape' in x and 'max_shape' in x: out['min_shape'] = x['min_shape'] out['max_shape'] = x['max_shape'] + if 'min_value' in x and 'max_value' in x: + np_dst_type = mstype.dtype_to_nptype(dst_type) + if isinstance(x['min_value'], (int, float, tuple, list)): + min_value = Tensor(np.array(x['min_value']).astype(np_dst_type)) + else: + min_value = Tensor(x['min_value'].asnumpy().astype(np_dst_type)) + min_value = tuple(min_value.asnumpy()) + if isinstance(x['max_value'], (int, float, tuple, list)): + max_value = Tensor(np.array(x['max_value']).astype(np_dst_type)) + else: + max_value = Tensor(x['max_value'].asnumpy().astype(np_dst_type)) + max_value = tuple(max_value.asnumpy()) + out['min_value'] = min_value + out['max_value'] = max_value return out @@ -892,7 +906,8 @@ class Gather(Primitive): >>> output = ops.Gather()(input_params, input_indices, axis) >>> print(output) [1. 3. 5. 3. 7.] - >>> # case2: input_indices is a Tensor with shape (2, 2). When the input_params has one dimension, the output shape is equal to the input_indices shape. + >>> # case2: input_indices is a Tensor with shape (2, 2). When the input_params has one dimension, + the output shape is equal to the input_indices shape. >>> input_indices = Tensor(np.array([[0, 2], [2, 6]]), mindspore.int32) >>> axis = 0 >>> output = ops.Gather()(input_params, input_indices, axis) @@ -2043,29 +2058,19 @@ class Tile(PrimitiveWithInfer): return (True, base_tensor) return (False, None) - def __infer__(self, x, multiples): - multiples_v = multiples['value'] - if multiples_v is None: - if len(multiples['shape']) != 1: - raise ValueError(f'For \'{self.name}\' the dim of multiples must be 1.') - rank = max(len(x['shape']), multiples['shape'][0]) - out_shape = [-1] * rank - # tile can't infer min/max shape if multiples_v is None - return {'shape': out_shape, - 'dtype': x['dtype'], - 'value': None, - 'min_shape': [1] * rank, - 'max_shape': [1] * rank - } - + def _get_shape_and_range(self, x, multiples): + """calculate tile shape and value""" x_shp = x['shape'] - validator.check_value_type( - "multiples", multiples_v, [tuple], self.name) - for i, multiple in enumerate(multiples_v): - validator.check_positive_int( - multiple, "multiples[%d]" % i, self.name) - validator.check_value_type( - "x[\'dtype\']", x["dtype"], mstype.tensor_type, self.name) + multiples_v = multiples['value'] + value = None + if multiples_v is None: + multiples_v = multiples['min_value'] + if 'max_shape' in x and 'min_shape' in x: + max_shape = x['max_shape'] + min_shape = x['min_shape'] + else: + max_shape = list(x_shp) + min_shape = list(x_shp) len_sub = len(multiples_v) - len(x_shp) multiples_w = None if len_sub == 0: @@ -2073,19 +2078,85 @@ class Tile(PrimitiveWithInfer): if len_sub > 0: for i in range(0, len_sub): x_shp.insert(0, 1) + min_shape.insert(0, 1) + max_shape.insert(0, 1) multiples_w = multiples_v elif len_sub < 0: raise ValueError(f"For '{self.name}', the length of 'multiples' can not be smaller than " f"the dimension of 'input_x', but got length of 'multiples': {len(multiples_v)} " f"and dimension of 'input_x': {len(x_shp)}.") - for i, a in enumerate(multiples_w): - x_shp[i] *= a - value = None - if x['value'] is not None: - value = Tensor(np.tile(x['value'].asnumpy(), multiples_w)) - return {'shape': x_shp, + if 'max_value' in multiples and 'min_value' in multiples: + multiples_v_max = multiples['max_value'] + multiples_v_min = multiples['min_value'] + i = 0 + for a, b in zip(multiples_v_min, multiples_v_max): + if isinstance(a, (Tensor_, Tensor)): + a = a.asnumpy() + if isinstance(b, (Tensor_, Tensor)): + b = b.asnumpy() + x_shp[i] *= a + if a != b: + x_shp[i] = -1 + min_shape[i] *= a + max_shape[i] *= b + i += 1 + else: + for i, a in enumerate(multiples_w): + x_shp[i] *= a + max_shape[i] *= a + min_shape[i] *= a + if x['value'] is not None: + value = Tensor(np.tile(x['value'].asnumpy(), multiples_w)) + out_shape = { + 'shape': x_shp, + 'max_shape': max_shape, + 'min_shape': min_shape + } + return out_shape, value + + def __infer__(self, x, multiples): + multiples_v = multiples['value'] + if multiples_v is None: + if 'max_value' not in multiples or 'min_value' not in multiples: + if len(multiples['shape']) != 1: + raise ValueError(f'For \'{self.name}\', the dim of multiples must be 1.') + rank = max(len(x['shape']), multiples['shape'][0]) + out_shape = [-1] * rank + return { + 'shape': out_shape, + 'dtype': x['dtype'], + 'value': None, + 'max_shape': [1] * rank, + 'min_shape': [1] * rank + } + out_shape, value = self._get_shape_and_range(x, multiples) + max_shape = out_shape.get('max_shape', None) + min_shape = out_shape.get('min_shape', None) + shape = out_shape.get('shape', None) + return { + 'shape': shape, 'dtype': x['dtype'], - 'value': value} + 'value': value, + 'max_shape': max_shape, + 'min_shape': min_shape + } + + validator.check_value_type( + "multiples", multiples_v, [tuple], self.name) + for i, multiple in enumerate(multiples_v): + validator.check_positive_int( + multiple, "multiples[%d]" % i, self.name) + validator.check_value_type( + "x[\'dtype\']", x["dtype"], mstype.tensor_type, self.name) + out_shp, value = self._get_shape_and_range(x, multiples) + shp = out_shp.get('shape', None) + out = {'shape': shp, + 'dtype': x['dtype'], + 'value': value} + if 'max_shape' in x and 'min_shape' in x: + out['max_shape'] = out_shp.get('max_shape', None) + out['min_shape'] = out_shp.get('min_shape', None) + return out class UnsortedSegmentSum(PrimitiveWithInfer): @@ -4099,6 +4170,28 @@ class TensorScatterUpdate(PrimitiveWithInfer): def __init__(self): self.init_prim_io_names(inputs=['input_x', 'indices', 'updates'], outputs=['y']) + def _infer_min_max_value(self, input_x_value, indices_value, updates_value): + """TensorScatterUpdate infer min max value""" + if isinstance(input_x_value, tuple): + input_x_value = list(input_x_value) + if isinstance(input_x_value, (Tensor, Tensor_)): + input_x_value = input_x_value.asnumpy() + indices = indices_value.asnumpy() + input_x = np.array(input_x_value) + updates = np.array(updates_value) + for i, indice in enumerate(indices): + input_x[indice] = updates[i] + output = tuple(input_x.tolist()) + return output + + def infer_min_value(self, input_x_value, indices_value, updates_value): + """TensorScatterUpdate infer min value""" + return self._infer_min_max_value(input_x_value, indices_value, updates_value) + + def infer_max_value(self, input_x_value, indices_value, updates_value): + """TensorScatterUpdate infer max value""" + return self._infer_min_max_value(input_x_value, indices_value, updates_value) + def infer_shape(self, input_x_shape, indices_shape, updates_shape): if len(indices_shape) < 2: raise ValueError(f"For '{self.name}', the dimension of 'indices' cannot be less than 2," diff --git a/mindspore/python/mindspore/ops/operations/math_ops.py b/mindspore/python/mindspore/ops/operations/math_ops.py index 0705648585d..2a7576ce812 100644 --- a/mindspore/python/mindspore/ops/operations/math_ops.py +++ b/mindspore/python/mindspore/ops/operations/math_ops.py @@ -25,6 +25,7 @@ from ...common.tensor import Tensor from ...common._decorator import deprecated from .._utils import get_broadcast_shape from ..primitive import Primitive, PrimitiveWithInfer, PrimitiveWithCheck, prim_attr_register, _run_op +from ..._c_expression import Tensor as Tensor_ def _infer_shape_reduce(x, axis, keep_dims, prim_name): @@ -234,6 +235,27 @@ class Add(_MathBinaryOp): >>> print(output.dtype) Float32 """ + def _add_infer_special_value(self, a, b): + """Add infer min max value""" + if a is not None and b is not None: + if isinstance(a, (Tensor, Tensor_)): + a = a.asnumpy() + if isinstance(b, (Tensor, Tensor_)): + b = b.asnumpy() + a = np.array(a) + b = np.array(b) + out = a + b + out = tuple(out.tolist()) + return out + return None + + def infer_min_value(self, x, y): + """Add infer min value""" + return _add_infer_special_value(x, y) + + def infer_max_value(self, x, y): + """Add infer min value""" + return _add_infer_special_value(x, y) def infer_value(self, x, y): if x is not None and y is not None: @@ -1915,6 +1937,27 @@ class Mul(_MathBinaryOp): >>> print(output) [ 4. 10. 18.] """ + def _infer_min_max_value(self, x, y): + """Mul infer min max value""" + if x is not None and y is not None: + if isinstance(x, (Tensor, Tensor_)): + x = x.asnumpy() + if isinstance(y, (Tensor, Tensor_)): + y = y.asnumpy() + x = np.array(x) + y = np.array(y) + out = x * y + out = tuple(out.tolist()) + return out + return None + + def infer_min_value(self, x, y): + """Mul infer min value""" + return self._infer_min_max_value(x, y) + + def infer_max_value(self, x, y): + """Mul infer max value""" + return self._infer_min_max_value(x, y) def infer_dtype(self, x_dtype, y_dtype): mul_valid_type = mstype.number_type + (mstype.bool_,) @@ -2829,6 +2872,27 @@ class Div(_MathBinaryOp): >>> print(output.dtype) Float32 """ + def _div_infer_special_value(self, x, y): + """Div infer min max value""" + if x is not None and y is not None: + if isinstance(x, (Tensor, Tensor_)): + x = x.asnumpy() + if isinstance(y, (Tensor, Tensor_)): + y = y.asnumpy() + x = np.array(x) + y = np.array(y) + out = x / y + out = tuple(out.tolist()) + return out + return None + + def infer_min_value(self, x, y): + """Div infer min value""" + return _div_infer_special_value(x, y) + + def infer_max_value(self, x, y): + """Div infer max value""" + return _div_infer_special_value(x, y) def infer_value(self, x, y): if x is not None and y is not None: diff --git a/mindspore/python/mindspore/ops/primitive.py b/mindspore/python/mindspore/ops/primitive.py index 7af70e401e5..fe49edca37d 100644 --- a/mindspore/python/mindspore/ops/primitive.py +++ b/mindspore/python/mindspore/ops/primitive.py @@ -576,7 +576,6 @@ class PrimitiveWithInfer(Primitive): # in non-graph_mode, it is not necessary to infer min/max shape if not is_graph_mode: return out - # output does not contain dynamic shape, no need to calculate min/max shape def has_dynamic_shape(shp): if isinstance(shp, int): @@ -585,6 +584,27 @@ class PrimitiveWithInfer(Primitive): return any(has_dynamic_shape(e) for e in shp) return False + # calculate min/max value for output + def get_specified_value(elems, attr): + has_specified_value = False + ret_vals = [] + for elem in elems: + if attr in elem: + has_specified_value = True + ret_vals.append(elem[attr]) + else: + ret_vals.append(elem['value']) + return has_specified_value, tuple(ret_vals) + + has_min_value, min_values = get_specified_value(args, 'min_value') + has_max_value, max_values = get_specified_value(args, 'max_value') + if has_min_value and has_max_value: + if hasattr(self, 'infer_min_value'): + fn_infer_min_value = getattr(self, 'infer_min_value') + out['min_value'] = fn_infer_min_value(*min_values) + if hasattr(self, 'infer_max_value'): + fn_infer_max_value = getattr(self, 'infer_max_value') + out['max_value'] = fn_infer_max_value(*max_values) if not has_dynamic_shape(out['shape']): return out