!47321 setitem-by-bool support broadcast

Merge pull request !47321 from chenweifeng/setitem-by-bool-broadcast

mindspore/python/mindspore/ops/composite/multitype_ops/_compile_utils.py

@@ -838,7 +838,8 @@ def _tensor_setitem_by_bool_tensor_with_tensor(data, index, value):
     """Set a tensor item by a bool tensor with a tensor."""
     dtype = F.dtype(data)
     u_cast = F.cast(value, dtype)
-    result = data * index + u_cast * F.logical_not(index)
+    index = index.reshape(const_utils.generate_padding_shape(index.shape, len(data.shape)))
+    result = u_cast * index + data * F.logical_not(index)
     return result
 
 

mindspore/python/mindspore/ops/composite/multitype_ops/_constexpr_utils.py

@@ -1025,3 +1025,14 @@ def promote_binary_dtype(dtype_1, dtype_2):
     if dtype_1 in complex_types or dtype_2 in complex_types:
         return get_output_dtype(dtype_1, dtype_2, True)
     return get_output_dtype(dtype_1, dtype_2, False)
+
+
+@constexpr
+def generate_padding_shape(shape, length):
+    """
+    pad the `shape` to `length` with 1.
+    """
+    if len(shape) > length:
+        raise ValueError(f"Can not pad {shape} to length {length}.")
+
+    return shape + (1,) * (length - len(shape))

tests/st/ops/ascend/test_tensor_setitem.py

@@ -0,0 +1,40 @@
+# 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.
+# ============================================================================
+""" test_tensor_setitem """
+import pytest
+import numpy as np
+from mindspore import Tensor
+
+
+@pytest.mark.level0
+@pytest.mark.platform_arm_ascend_training
+@pytest.mark.platform_x86_ascend_training
+@pytest.mark.env_onecard
+def test_tensor_slice_by_bool_broadcast():
+    """
+    Feature: Tensor-setitem-by-bool support broadcast.
+    Description: Tensor-setitem-by-bool support broadcast.
+    Expectation: success.
+    """
+    data_np = np.ones([2, 3, 4], np.float32)
+    index_np = np.array([True, False])
+    value = 2
+
+    data_tensor = Tensor(data_np)
+    index_tensor = Tensor(index_np)
+
+    data_np[index_np] = value
+    data_tensor[index_tensor] = value
+    assert np.allclose(data_tensor.asnumpy(), data_np)