Fix gpu and add autograd test

src/operator/contrib/boolean_mask.cc

@@ -154,7 +154,7 @@ inline void BooleanMaskBackward<cpu>(const nnvm::NodeAttrs& attrs,
   const NDArray& idx = inputs[2];
   const NDArray& igrad_data = outputs[0];
   MSHADOW_TYPE_SWITCH(igrad_data.dtype(), DType, {
-    MSHADOW_TYPE_SWITCH(idx.dtype(), IType, {
+    MSHADOW_TYPE_SWITCH_WITH_BOOL(idx.dtype(), IType, {
       size_t input_size = igrad_data.shape().Size();
       size_t idx_size = idx.shape()[0];
       size_t col_size = input_size / idx_size;

src/operator/contrib/boolean_mask.cu

@@ -66,7 +66,7 @@ inline void BooleanMaskForward<gpu>(const nnvm::NodeAttrs& attrs,
     ctx.requested[0].get_space_typed<gpu, 1, char>(Shape1(temp_storage_bytes), s);
   prefix_sum = reinterpret_cast<int32_t*>(workspace.dptr_);
   d_temp_storage = workspace.dptr_ + buffer_size;
-  MSHADOW_TYPE_SWITCH(idx.dtype(), IType, {
+  MSHADOW_TYPE_SWITCH_WITH_BOOL(idx.dtype(), IType, {
     mxnet_op::Kernel<mshadow_op::identity_with_cast, gpu>::Launch(
       s, idx.shape()[0], prefix_sum, idx.data().dptr<IType>());
   });
@@ -129,7 +129,7 @@ inline void BooleanMaskBackward<gpu>(const nnvm::NodeAttrs& attrs,
     ctx.requested[0].get_space_typed<gpu, 1, char>(Shape1(temp_storage_bytes), s);
   prefix_sum = reinterpret_cast<int32_t*>(workspace.dptr_);
   d_temp_storage = workspace.dptr_ + buffer_size;
-  MSHADOW_TYPE_SWITCH(idx.dtype(), IType, {
+  MSHADOW_TYPE_SWITCH_WITH_BOOL(idx.dtype(), IType, {
     mxnet_op::Kernel<mshadow_op::identity_with_cast, gpu>::Launch(
       s, idx.shape()[0], prefix_sum, idx.data().dptr<IType>());
   });

tests/python/unittest/test_numpy_ndarray.py

@@ -934,8 +934,8 @@ def test_np_multinomial():
 @with_seed()
 @use_np
 def test_np_ndarray_boolean_indexing():
-    # adapted from numpy's test_indexing.py
     def test_single_bool_index():
+        # adapted from numpy's test_indexing.py
         # Single boolean index
         a = np.array([[1, 2, 3],
                       [4, 5, 6],
@@ -944,6 +944,7 @@ def test_single_bool_index():
         assert same(a[np.array(False, dtype=np.bool_)].asnumpy(), a[None][0:0].asnumpy())
 
     def test_boolean_catch_exception():
+        # adapted from numpy's test_indexing.py
         arr = np.ones((5, 4, 3))
 
         index = np.array([True], dtype=np.bool_)
@@ -958,13 +959,15 @@ def test_boolean_catch_exception():
         assert_exception(arr.__getitem__, TypeError, (slice(None), index))
 
     def test_boolean_indexing_onedim():
+        # adapted from numpy's test_indexing.py
         # Indexing a 2-dimensional array with
         # boolean array of length one
         a = np.array([[0.,  0.,  0.]])
         b = np.array([True], dtype=bool)
         assert same(a[b].asnumpy(), a.asnumpy())
 
     def test_boolean_indexing_twodim():
+        # adapted from numpy's test_indexing.py
         # Indexing a 2-dimensional array with
         # 2-dimensional boolean array
         a = np.array([[1, 2, 3],
@@ -978,17 +981,34 @@ def test_boolean_indexing_twodim():
         assert same(a[b[0]].asnumpy(), a[b[2]].asnumpy())
 
     def test_boolean_indexing_list():
+        # adapted from numpy's test_indexing.py
         a = np.array([1, 2, 3], dtype=np.int32)
         b = [True, False, True]
         # Two variants of the test because the first takes a fast path
         assert same(a[b].asnumpy(), _np.array([1, 3], dtype=a.dtype))
         (a[None, b], [[1, 3]])
 
+    def test_boolean_indexing_autograd():
+        a = np.random.uniform(size=(3, 4, 5))
+        a.attach_grad()
+        with mx.autograd.record():
+            out_mx = a[a < 0.5]
+        out_mx.backward()
+
+        a_np = a.asnumpy()
+        out_np = a_np[a_np < 0.5]
+        assert_almost_equal(out_mx.asnumpy(), out_np, rtol=1e-4, atol=1e-5, use_broadcast=False)
+
+        a_grad_np = _np.zeros(a.shape, dtype=a.dtype)
+        a_grad_np[a_np < 0.5] = 1
+        assert_almost_equal(a.grad.asnumpy(), a_grad_np, rtol=1e-4, atol=1e-5, use_broadcast=False)
+
     test_single_bool_index()
     test_boolean_catch_exception()
     test_boolean_indexing_onedim()
     test_boolean_indexing_twodim()
     test_boolean_indexing_list()
+    test_boolean_indexing_autograd()
 
 
 if __name__ == '__main__':