From 9cad0fec1d717b935590a8ab74bf2b4f7263c639 Mon Sep 17 00:00:00 2001 From: wilfChen Date: Wed, 29 Jul 2020 16:24:01 +0800 Subject: [PATCH] gpu broadcast to --- .../kernel_compiler/gpu/arrays/broadcast_to_gpu_kernel.h | 7 ++++--- tests/st/ops/gpu/test_broadcast_to_ops.py | 6 ++++++ 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/broadcast_to_gpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/broadcast_to_gpu_kernel.h index 1bbaff80f6..280879b81c 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/broadcast_to_gpu_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/broadcast_to_gpu_kernel.h @@ -51,11 +51,12 @@ class BroadcastToGpuKernel : public GpuKernel { MS_LOG(EXCEPTION) << "BroadcastTo operation not support dim greater than 4"; } - for (int i = input_shapes.size() - 1; i >= 0; i--) { - input_shape_[i] = input_shapes[i]; + size_t offset = output_shapes.size() - input_shapes.size(); + for (size_t i = 0; i < input_shapes.size(); i++) { + input_shape_[i + offset] = input_shapes[i]; } - for (int j = output_shapes.size() - 1; j >= 0; j--) { + for (size_t j = 0; j < output_shapes.size(); j++) { output_shape_[j] = output_shapes[j]; } diff --git a/tests/st/ops/gpu/test_broadcast_to_ops.py b/tests/st/ops/gpu/test_broadcast_to_ops.py index 828e72c4d0..137f271519 100644 --- a/tests/st/ops/gpu/test_broadcast_to_ops.py +++ b/tests/st/ops/gpu/test_broadcast_to_ops.py @@ -38,3 +38,9 @@ def test_broadcast(): output = P.BroadcastTo(shape)(Tensor(x1_np)) expect = np.broadcast_to(x1_np, shape) assert np.allclose(output.asnumpy(), expect) + + x1_np = np.random.rand(4, 5).astype(np.float32) + shape = (2, 3, 4, 5) + output = P.BroadcastTo(shape)(Tensor(x1_np)) + expect = np.broadcast_to(x1_np, shape) + assert np.allclose(output.asnumpy(), expect)