This is an automated email from the ASF dual-hosted git repository.
wkcn 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 8c76631 Relaxing type requirements for broadcast_like (#17977)
8c76631 is described below
commit 8c76631caa5f349aef5d515ec95a64d2954dbcef
Author: wicky <tbc.dengwe...@gmail.com>
AuthorDate: Mon May 4 14:03:33 2020 +0800
Relaxing type requirements for broadcast_like (#17977)
* Relaxing type requirements for broadcast_like
* enhance unit test
---
src/operator/tensor/broadcast_reduce_op_value.cc | 11 ++++++++++-
tests/python/unittest/test_operator.py | 10 ++++++++++
2 files changed, 20 insertions(+), 1 deletion(-)
diff --git a/src/operator/tensor/broadcast_reduce_op_value.cc
b/src/operator/tensor/broadcast_reduce_op_value.cc
index 0a14a20..71be8f8 100644
--- a/src/operator/tensor/broadcast_reduce_op_value.cc
+++ b/src/operator/tensor/broadcast_reduce_op_value.cc
@@ -138,7 +138,16 @@ NNVM_REGISTER_OP(broadcast_like)
[](const NodeAttrs& attrs) {
return std::vector<std::string>{"lhs", "rhs"};
})
-.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<2, 1>)
+.set_attr<nnvm::FInferType>("FInferType", [](const nnvm::NodeAttrs& attrs,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ CHECK_EQ(in_attrs->size(), 2) << " in operator " << attrs.name;
+ std::vector<int> checked_in_attrs = { (*in_attrs)[0] };
+ bool ret = !type_is_none((*in_attrs)[1]) &&
+ ElemwiseType<1, 1>(attrs, &checked_in_attrs, out_attrs);
+ (*in_attrs)[0] = checked_in_attrs[0];
+ return ret;
+})
.set_attr<nnvm::FGradient>("FGradient",
[](const nnvm::ObjectPtr& n,
const std::vector<nnvm::NodeEntry>& ograds) {
diff --git a/tests/python/unittest/test_operator.py
b/tests/python/unittest/test_operator.py
index 0baa941..4b8ddd6 100644
--- a/tests/python/unittest/test_operator.py
+++ b/tests/python/unittest/test_operator.py
@@ -3120,6 +3120,16 @@ def test_reshape_like_different_types():
assert_allclose(z.asnumpy(), [[0,0],[0,0],[0,0]])
@with_seed()
+def test_broadcast_like_different_types():
+ x = mx.nd.zeros((2, 1))
+ y = mx.nd.ones((2, 2))
+
+ y = mx.nd.array(y).astype('int32')
+ z = mx.nd.broadcast_like(x, y)
+ assert_allclose(z.asnumpy(), [[0,0],[0,0]])
+ assert x.dtype == z.dtype
+
+@with_seed()
def test_flip():
for ndim in range(1, 6):
for t in range(5):
