[incubator-mxnet] branch master updated: Fix a bug in where op with 1-D input (#12325)

[zh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zhasheng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git


The following commit(s) were added to refs/heads/master by this push:
     new 91730fd  Fix a bug in where op with 1-D input (#12325)
91730fd is described below

commit 91730fdb6180ddcc3ab171df72c37f0ac15b59fd
Author: Lin Yuan <ap...@gmail.com>
AuthorDate: Fri Aug 24 12:51:18 2018 -0700

    Fix a bug in where op with 1-D input (#12325)
    
    * Fix a bug in where op with 1-D input
    
    * Add unit test
---
 src/operator/tensor/control_flow_op.h  | 1 +
 tests/python/unittest/test_operator.py | 9 +++++++++
 2 files changed, 10 insertions(+)

diff --git a/src/operator/tensor/control_flow_op.h b/src/operator/tensor/control_flow_op.h
index 94e6510..e9aa9f6 100644
--- a/src/operator/tensor/control_flow_op.h
+++ b/src/operator/tensor/control_flow_op.h
@@ -189,6 +189,7 @@ inline bool WhereOpShape(const nnvm::NodeAttrs& attrs,
     return true;
   } else if ((*in_attrs)[0].ndim() == 1) {
     CHECK_EQ((*in_attrs)[0].Size(), static_cast<size_t>(tshape[0]));
+    return true;
   }
   return false;
 }
diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py
index fd60611..d0bc450 100644
--- a/tests/python/unittest/test_operator.py
+++ b/tests/python/unittest/test_operator.py
@@ -4515,6 +4515,14 @@ def test_where():
                                              y=mx.nd.array([[8,9],[10,11],[12,13]]),
                                              condition=mx.nd.array([1,0])), MXNetError)
 
+    def test_1d_cond():
+        cond = mx.nd.array([1, 0, 1])
+        x = mx.nd.array([[2, 3], [4, 5], [6, 7]])
+        y = mx.nd.array([[7, 8], [9, 10], [10, 11]])
+        expect_out = np.array([[2, 3], [9, 10], [6, 7]])
+        out = mx.nd.where(cond, x, y).asnumpy()
+        assert(expect_out.all() == out.all())
+
     test_where_helper((5, 9), True)
     test_where_helper((5, 9), False)
     test_where_helper((5, 7, 9), True)
@@ -4526,6 +4534,7 @@ def test_where():
     test_where_numeric_gradient((5, 7, 9), True)
     test_where_numeric_gradient((5, 7, 9), False)
     test_invalid_shape()
+    test_1d_cond()
 
 @with_seed()
 def test_new_softmax():