apeforest commented on a change in pull request #17882:
URL: https://github.com/apache/incubator-mxnet/pull/17882#discussion_r411540591
##########
File path: src/operator/tensor/broadcast_reduce_op.h
##########
@@ -1058,22 +1065,58 @@ struct broadcast_kernel {
mshadow::Shape<MXNET_SPECIAL_MAX_NDIM>
in_shape,
mshadow::Shape<MXNET_SPECIAL_MAX_NDIM>
out_shape,
const OpReqType req,
- const uint32_t ndim) {
- size_t in_stride = 1;
- size_t out_stride = 1;
- index_t idx = i;
- index_t in_idx = i;
- for (int iter = ndim - 1; iter >= 0; --iter) {
- size_t dim_idx = idx % out_shape[iter];
- in_idx -= dim_idx * out_stride;
- if (in_shape[iter] != 1) {
- in_idx += dim_idx * in_stride;
- }
- idx /= out_shape[iter];
- in_stride *= in_shape[iter];
- out_stride *= out_shape[iter];
- }
- KERNEL_ASSIGN(output[i], req, OP::Map(input[in_idx]));
+ const uint32_t ndim,
Review comment:
shall we use `int` for ndim as we do for all other places?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
