Relaxing type requirements for broadcast_like (apache#17977) (apache#…

…19447) * Relaxing type requirements for broadcast_like * enhance unit test Co-authored-by: wicky <[email protected]>
josephevans · Dec 8, 2020 · 5434d41 · 5434d41
1 parent 9964639
commit 5434d41
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/src/operator/tensor/broadcast_reduce_op_value.cc b/src/operator/tensor/broadcast_reduce_op_value.cc
@@ -137,7 +137,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
@@ -3165,6 +3165,16 @@ def test_reshape_like_different_types():
     z = mx.nd.reshape_like(x, y)
     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):