Merge remote-tracking branch 'upstream/master' into windows_builds

python/mxnet/ndarray/ndarray.py

@@ -963,7 +963,10 @@ def _set_nd_basic_indexing(self, key, value):
                     value_nd.copyto(self)
 
             elif isinstance(value, numeric_types):
-                self._full(value)
+                if isinstance(value, bool):
+                    self._full(int(value))
+                else:
+                    self._full(value)
 
             elif isinstance(value, (np.ndarray, np.generic)):
                 tmp_shape = _shape_for_bcast(

python/mxnet/numpy/multiarray.py

@@ -480,7 +480,9 @@ def _set_np_boolean_indexing(self, key, value):
                 pos -= 1
 
         if isinstance(value, numeric_types):
-            _npi.boolean_mask_assign_scalar(data=data, mask=mask, value=value, start_axis=pos, out=data)
+            _npi.boolean_mask_assign_scalar(data=data, mask=mask,
+                                            value=int(value) if isinstance(value, bool) else value,
+                                            start_axis=pos, out=data)
         elif isinstance(value, ndarray):
             _npi.boolean_mask_assign_tensor(data=data, mask=mask, value=value, start_axis=pos, out=data)
         else:

src/operator/numpy/np_boolean_mask_assign.cc

@@ -91,14 +91,21 @@ struct BooleanAssignCPUKernel {
                   const size_t middle,
                   const size_t valid_num,
                   const size_t trailing,
-                  DType* tensor) {
+                  DType* tensor,
+                  const bool broadcast = false) {
     // binary search for the turning point
     size_t mid = bin_search(idx, idx_size, i);
     // final answer is in mid
     for (size_t l = 0; l < leading; ++l) {
       for (size_t t = 0; t < trailing; ++t) {
-        data[(l * middle + mid) * trailing + t] =
-          (scalar) ? tensor[0] : tensor[(l * valid_num + i) * trailing + t];
+        if (scalar) {
+          data[(l * middle + mid) * trailing + t] = tensor[0];
+        } else {
+          data[(l * middle + mid) * trailing + t] =
+            (broadcast) ?
+            tensor[l * trailing + t] :
+            tensor[(l * valid_num + i) * trailing + t];
+        }
       }
     }
   }
@@ -134,19 +141,17 @@ bool BooleanAssignShape(const nnvm::NodeAttrs& attrs,
     } else {
       const TShape& vshape = in_attrs->at(2);
       if (vshape.Size() > 1) {
-        for (int i = 0; i < dshape.ndim(); ++i) {
-          if (i < start_axis) {
-            CHECK_EQ(dshape[i], vshape[i])
-              << "shape mismatch of value with input at dimension " << i
-              << "; dimension is " << dshape[i] << " but corresponding value dimension is "
-              << vshape[i];
-          }
-          if (i >= start_axis + mshape.ndim()) {
-            CHECK_EQ(dshape[i], vshape[i - mshape.ndim() + 1])
-              << "shape mismatch of value with input at dimension " << i
-              << "; dimension is " << dshape[i] << " but corresponding value dimension is "
-              << vshape[i - mshape.ndim() + 1];
-          }
+        for (int i = 0; i < start_axis; ++i) {
+          CHECK_EQ(dshape[i], vshape[i])
+            << "shape mismatch of value with input at dimension " << i
+            << "; dimension is " << dshape[i] << " but corresponding value dimension is "
+            << vshape[i];
+        }
+        for (int i = 1; i <= (dshape.ndim() - start_axis - mshape.ndim()); ++i) {
+          CHECK_EQ(dshape[dshape.ndim() - i], vshape[vshape.ndim() - i])
+            << "shape mismatch of value with input at dimension " << (dshape.ndim() - i)
+            << "; dimension is " << dshape[dshape.ndim() - i]
+            << " but corresponding value dimension is " << vshape[vshape.ndim() - 1];
         }
       }
     }
@@ -215,14 +220,22 @@ void NumpyBooleanAssignForwardCPU(const nnvm::NodeAttrs& attrs,
   // If there's no True in mask, return directly
   if (valid_num == 0) return;
 
+  const TShape& vshape = inputs[2].shape_;
+
   if (inputs.size() == 3U) {
-    const TShape& vshape = inputs[2].shape_;
+    // tensor case
     if (inputs[2].shape_.Size() != 1) {
-      // tensor case, check tensor size with the valid_num
-      CHECK_EQ(static_cast<size_t>(valid_num), vshape[start_axis])
-        << "boolean array indexing assignment cannot assign " << vshape
-        << " input values to the " << valid_num << " output values where the mask is true"
-        << std::endl;
+      auto vndim = vshape.ndim();
+      auto dndim = dshape.ndim();
+      auto mndim = mshape.ndim();
+      CHECK(vndim <= (dndim - mndim + 1));
+      if ((vndim == (dndim - mndim + 1)) && (vshape[start_axis] != 1)) {
+        // tensor case, check tensor size equal to or broadcastable with valid_num
+        CHECK_EQ(static_cast<size_t>(valid_num), vshape[start_axis])
+          << "boolean array indexing assignment cannot assign " << vshape
+          << " input values to the " << valid_num << " output values where the mask is true"
+          << std::endl;
+      }
     }
   }
 
@@ -240,20 +253,23 @@ void NumpyBooleanAssignForwardCPU(const nnvm::NodeAttrs& attrs,
   }
 
   if (inputs.size() == 3U) {
-    MSHADOW_TYPE_SWITCH(data.type_flag_, DType, {
+    MSHADOW_TYPE_SWITCH_WITH_BOOL(data.type_flag_, DType, {
       if (inputs[2].shape_.Size() == 1) {
         Kernel<BooleanAssignCPUKernel<true>, cpu>::Launch(
           s, valid_num, data.dptr<DType>(), prefix_sum.data(), prefix_sum.size(),
           leading, middle, valid_num, trailing, inputs[2].dptr<DType>());
       } else {
-       Kernel<BooleanAssignCPUKernel<false>, cpu>::Launch(
+        bool need_broadcast = (vshape.ndim() == (dshape.ndim() - mshape.ndim() + 1)) ?
+                              (vshape[start_axis] == 1) :
+                              true;
+        Kernel<BooleanAssignCPUKernel<false>, cpu>::Launch(
           s, valid_num, data.dptr<DType>(), prefix_sum.data(), prefix_sum.size(),
-          leading, middle, valid_num, trailing, inputs[2].dptr<DType>());
+          leading, middle, valid_num, trailing, inputs[2].dptr<DType>(), need_broadcast);
       }
     });
   } else {
     CHECK(attrs.dict.find("value") != attrs.dict.end()) << "value needs be provided";
-    MSHADOW_TYPE_SWITCH(data.type_flag_, DType, {
+    MSHADOW_TYPE_SWITCH_WITH_BOOL(data.type_flag_, DType, {
       Kernel<BooleanAssignCPUKernel<true>, cpu>::Launch(
         s, valid_num, data.dptr<DType>(), prefix_sum.data(), prefix_sum.size(),
         leading, middle, trailing, static_cast<DType>(std::stod(attrs.dict.at("value"))));

src/operator/numpy/np_boolean_mask_assign.cu

@@ -93,7 +93,8 @@ struct BooleanAssignGPUKernel {
                              const size_t middle,
                              const size_t valid_num,
                              const size_t trailing,
-                             DType* tensor) {
+                             DType* tensor,
+                             const bool broadcast = false) {
     // binary search for the turning point
     size_t m = i / trailing % valid_num;
     size_t l = i / trailing / valid_num;
@@ -103,7 +104,7 @@ struct BooleanAssignGPUKernel {
     if (scalar) {
       data[dst] = tensor[0];
     } else {
-      data[dst] = tensor[i];
+      data[dst] = broadcast ? tensor[l * trailing + i % trailing] : tensor[i];
     }
   }
 };
@@ -200,14 +201,22 @@ void NumpyBooleanAssignForwardGPU(const nnvm::NodeAttrs& attrs,
   // If there's no True in mask, return directly
   if (valid_num == 0) return;
 
+  const TShape& vshape = inputs[2].shape_;
+
   if (inputs.size() == 3U) {
-    const TShape& vshape = inputs[2].shape_;
+    // tensor case
     if (inputs[2].shape_.Size() != 1) {
-      // tensor case, check tensor size with the valid_num
-      CHECK_EQ(static_cast<size_t>(valid_num), vshape[start_axis])
-        << "boolean array indexing assignment cannot assign " << vshape
-        << " input values to the " << valid_num << " output values where the mask is true"
-        << std::endl;
+      auto vndim = vshape.ndim();
+      auto dndim = dshape.ndim();
+      auto mndim = mshape.ndim();
+      CHECK(vndim <= (dndim - mndim + 1));
+      if ((vndim == (dndim - mndim + 1)) && (vshape[start_axis] != 1)) {
+        // tensor case, check tensor size equal to or broadcastable with valid_num
+        CHECK_EQ(static_cast<size_t>(valid_num), vshape[start_axis])
+          << "boolean array indexing assignment cannot assign " << vshape
+          << " input values to the " << valid_num << " output values where the mask is true"
+          << std::endl;
+      }
     }
   }
 
@@ -226,22 +235,25 @@ void NumpyBooleanAssignForwardGPU(const nnvm::NodeAttrs& attrs,
 
   if (inputs.size() == 3U) {
     if (inputs[2].shape_.Size() == 1) {
-      MSHADOW_TYPE_SWITCH(data.type_flag_, DType, {
+      MSHADOW_TYPE_SWITCH_WITH_BOOL(data.type_flag_, DType, {
         Kernel<BooleanAssignGPUKernel<true>, gpu>::Launch(
           s, leading * valid_num * trailing, data.dptr<DType>(), prefix_sum, mask_size + 1,
           leading, middle, valid_num, trailing, inputs[2].dptr<DType>());
       });
     } else {
-      MSHADOW_TYPE_SWITCH(data.type_flag_, DType, {
+      bool need_broadcast = (vshape.ndim() == (dshape.ndim() - mshape.ndim() + 1)) ?
+                            (vshape[start_axis] == 1) :
+                            true;
+      MSHADOW_TYPE_SWITCH_WITH_BOOL(data.type_flag_, DType, {
         Kernel<BooleanAssignGPUKernel<false>, gpu>::Launch(
           s, leading * valid_num * trailing, data.dptr<DType>(), prefix_sum, mask_size + 1,
-          leading, middle, valid_num, trailing, inputs[2].dptr<DType>());
+          leading, middle, valid_num, trailing, inputs[2].dptr<DType>(), need_broadcast);
       });
     }
   } else {
     CHECK(attrs.dict.find("value") != attrs.dict.end()) << "value is not provided";
     double value = std::stod(attrs.dict.at("value"));
-    MSHADOW_TYPE_SWITCH(data.type_flag_, DType, {
+    MSHADOW_TYPE_SWITCH_WITH_BOOL(data.type_flag_, DType, {
       Kernel<BooleanAssignGPUKernel<true>, gpu>::Launch(
         s, leading * valid_num * trailing, data.dptr<DType>(), prefix_sum, mask_size + 1,
         leading, middle, valid_num, trailing, static_cast<DType>(value));

src/operator/tensor/indexing_op.cc

@@ -296,8 +296,8 @@ void TakeOpForward<cpu>(const nnvm::NodeAttrs& attrs,
   Stream<cpu> *s = ctx.get_stream<cpu>();
   const int actual_axis = param.axis + ((param.axis < 0) ? arrshape.ndim() : 0);
 
-  MSHADOW_TYPE_SWITCH(outputs[take_::kOut].type_flag_, DType, {  // output data type
-    MSHADOW_TYPE_SWITCH(inputs[take_::kIdx].type_flag_, IType, {  // index data type
+  MSHADOW_TYPE_SWITCH_WITH_BOOL(outputs[take_::kOut].type_flag_, DType, {  // output data type
+    MSHADOW_TYPE_SWITCH_WITH_BOOL(inputs[take_::kIdx].type_flag_, IType, {  // index data type
       if (param.mode == take_::kRaise) {
         IType min = 0;
         IType max = static_cast<IType>(arrshape[actual_axis] - 1);

src/operator/tensor/indexing_op.cu

@@ -550,8 +550,8 @@ void TakeOpForward<gpu>(const nnvm::NodeAttrs& attrs,
   Stream<gpu> *s = ctx.get_stream<gpu>();
   const int actual_axis = param.axis + ((param.axis < 0) ? arrshape.ndim() : 0);
 
-  MSHADOW_TYPE_SWITCH(outputs[take_::kOut].type_flag_, DType, {  // output data type
-    MSHADOW_TYPE_SWITCH(inputs[take_::kIdx].type_flag_, IType, {  // index data type
+  MSHADOW_TYPE_SWITCH_WITH_BOOL(outputs[take_::kOut].type_flag_, DType, {  // output data type
+    MSHADOW_TYPE_SWITCH_WITH_BOOL(inputs[take_::kIdx].type_flag_, IType, {  // index data type
       if (param.mode == take_::kRaise) {
         // check out-of-bound indices
         IType min = 0;

src/operator/tensor/init_op.h

@@ -252,7 +252,7 @@ struct InitOpWithScalarParam : dmlc::Parameter<InitOpWithScalarParam> {
       .describe("Context of output, in format [cpu|gpu|cpu_pinned](n)."
                   "Only used for imperative calls.");
     DMLC_DECLARE_FIELD(dtype).set_default(mshadow::kFloat32)
-      MXNET_ADD_ALL_TYPES
+      MXNET_ADD_ALL_TYPES_WITH_BOOL
       .describe("Target data type.");
     DMLC_DECLARE_FIELD(value)
       .describe("Value with which to fill newly created tensor");

tests/python/unittest/test_numpy_op.py

@@ -1435,22 +1435,29 @@ def hybrid_forward(self, F, a, mask, value):
         for config in configs:
             dshape, mshape, start_axis = config
             test_data = np.random.uniform(size=dshape)
-            np_mask = _np.random.choice(a=[False, True], size=mshape)
-            mx_mask = np.array(np_mask)
-            # to avoid using tvm op greater_scalar in gpu environment
-            valid_num = int(mx_mask.asnumpy().sum())
+            valid_num = 0
+            while valid_num == 0:
+                mx_mask = np.random.choice(np.array([False, True], dtype=np.bool), size=mshape)
+                if test_data.size == 0:
+                    break
+                valid_num = int(mx_mask.asnumpy().sum())
+            np_mask = mx_mask.asnumpy().astype(_np.bool)
             vshape = []
+            vshape_broadcast = []
             for i in range(len(dshape)):
                 if i < start_axis:
                     vshape.append(dshape[i])
+                    vshape_broadcast.append(dshape[i])
                 elif i == start_axis:
                     vshape.append(valid_num)
+                    vshape_broadcast.append(1)
                 elif i >= start_axis + len(mshape):
                     vshape.append(dshape[i])
-            vshape = tuple(vshape)
-            for val in [42.0, _np.array(42.), _np.array([42.]), _np.random.uniform(size=vshape)]:
+                    vshape_broadcast.append(dshape[i])
+            vshape_broadcast = tuple(vshape_broadcast)
+            for val in [42.0, _np.array(42.), _np.array([42.]), _np.random.uniform(size=vshape), _np.random.uniform(size=vshape_broadcast)]:
                 mx_val = val if isinstance(val, float) else np.array(val, dtype=np.float32)
-                test_block = TestBooleanAssignScalar(mx_val, start_axis) if isinstance(mx_val, float) else TestBooleanAssignTensor(start_axis)
+                test_block = TestBooleanAssignScalar(val, start_axis) if isinstance(val, float) else TestBooleanAssignTensor(start_axis)
                 if hybridize:
                     test_block.hybridize()
                 np_data = test_data.asnumpy()