diff --git a/mindspore/python/mindspore/common/tensor.py b/mindspore/python/mindspore/common/tensor.py index 6c075b1c3d2..e758f798f93 100644 --- a/mindspore/python/mindspore/common/tensor.py +++ b/mindspore/python/mindspore/common/tensor.py @@ -2699,6 +2699,8 @@ class Tensor(Tensor_): [0. 1. 0. 0.] """ self._init_check() + if isinstance(value, (float, int)): + value = tensor_operator_registry.get("scalar_to_tensor")(value, self.dtype) if not isinstance(mask, Tensor): raise TypeError("For 'Tensor.masked_fill', the type of the argument 'mask' must be Tensor, but " "got {}.".format(type(mask))) diff --git a/mindspore/python/mindspore/nn/layer/pooling.py b/mindspore/python/mindspore/nn/layer/pooling.py index 51fb44b9012..4ad9b66f570 100644 --- a/mindspore/python/mindspore/nn/layer/pooling.py +++ b/mindspore/python/mindspore/nn/layer/pooling.py @@ -753,7 +753,6 @@ class AdaptiveMaxPool3d(Cell): `ouput_size` can be a tuple with 3 elements, or a single D for :math:`(D, D, D)`. :math:`D`, :math:`H` and :math:`W` can be int or None which means the output size is the same as that of the input. - return_indices (bool): If `return_indices` is True, the indices of max value would be output. Default: False. diff --git a/mindspore/python/mindspore/ops/function/array_func.py b/mindspore/python/mindspore/ops/function/array_func.py index d26608e9108..63395f67ec6 100644 --- a/mindspore/python/mindspore/ops/function/array_func.py +++ b/mindspore/python/mindspore/ops/function/array_func.py @@ -71,7 +71,6 @@ tensor_scatter_max_ = P.TensorScatterMax() scalar_to_array_ = P.ScalarToArray() scalar_to_tensor_ = P.ScalarToTensor() tuple_to_array_ = P.TupleToArray() -masked_fill_ = P.MaskedFill() masked_select_ = P.MaskedSelect() matrix_band_part_ = P.array_ops.MatrixBandPart() ger_ = P.Ger() @@ -3494,6 +3493,9 @@ def masked_fill(input_x, mask, value): >>> print(output) [0.5 0.5 3. 0.5] """ + if isinstance(value, (float, int)) and isinstance(input_x, Tensor): + value = scalar_to_tensor_(value, input_x.dtype) + masked_fill_ = _get_cache_prim(P.MaskedFill)() return masked_fill_(input_x, mask, value) diff --git a/mindspore/python/mindspore/ops/functional.py b/mindspore/python/mindspore/ops/functional.py index 4d13097af47..9c39b882aca 100644 --- a/mindspore/python/mindspore/ops/functional.py +++ b/mindspore/python/mindspore/ops/functional.py @@ -834,6 +834,7 @@ tensor_operator_registry.register('sum', P.ReduceSum) tensor_operator_registry.register('split', P.Split) tensor_operator_registry.register('select', P.Select) tensor_operator_registry.register('zeros_like', P.ZerosLike) +tensor_operator_registry.register('scalar_to_tensor', scalar_to_tensor) tensor_operator_registry.register('masked_fill', masked_fill) tensor_operator_registry.register('masked_select', masked_select) tensor_operator_registry.register('nonzero', nonzero) diff --git a/tests/st/ops/gpu/test_maskedfill_op.py b/tests/st/ops/gpu/test_maskedfill_op.py index 27da1042408..64055fd853a 100644 --- a/tests/st/ops/gpu/test_maskedfill_op.py +++ b/tests/st/ops/gpu/test_maskedfill_op.py @@ -20,6 +20,7 @@ import mindspore.context as context import mindspore.nn as nn from mindspore import Tensor from mindspore.ops import operations as P +from mindspore.ops import functional as F context.set_context(mode=context.GRAPH_MODE, device_target="GPU") @@ -131,6 +132,40 @@ def test_maskedfill_float_value(): maskedfill_value(0.5) +@pytest.mark.level0 +@pytest.mark.platform_x86_gpu_training +@pytest.mark.env_onecard +def test_func_masked_fill_float(): + """ + Feature: Test func masked_fill. + Description: Test func masked_fill api with float value. + Expectation: The result match to expect. + """ + inputs = Tensor(np.array([[1, 2, 3, 4], [5, 6, 7, 8]]).astype(np.float16)) + mask = Tensor(np.array([[True, True, False, True], [False, False, True, False]]).astype(np.bool)) + value = 22 + expect = np.array([[22, 22, 3, 22], [5, 6, 22, 8]]).astype(np.float16) + output = F.masked_fill(inputs, mask, value) + assert (output.asnumpy() == expect).all() + + +@pytest.mark.level0 +@pytest.mark.platform_x86_gpu_training +@pytest.mark.env_onecard +def test_tensor_masked_fill_float(): + """ + Feature: Test Tensor masked_fill. + Description: Test Tensor masked_fill api with float value. + Expectation: The result match to expect. + """ + inputs = Tensor(np.array([[1, 2, 3, 4], [5, 6, 7, 8]]).astype(np.float16)) + mask = Tensor(np.array([[True, True, False, True], [False, False, True, False]]).astype(np.bool)) + value = 22 + output = inputs.masked_fill(mask, value) + expect = np.array([[22, 22, 3, 22], [5, 6, 22, 8]]).astype(np.float16) + assert (output.asnumpy() == expect).all() + + @pytest.mark.level0 @pytest.mark.platform_x86_gpu_training @pytest.mark.env_onecard