From 5434d416eea13c31e0f71dd5fccec492c3ae5151 Mon Sep 17 00:00:00 2001
From: Sam Skalicky <samskalicky@gmail.com>
Date: Thu, 29 Oct 2020 14:50:56 -0700
Subject: [PATCH] Relaxing type requirements for broadcast_like (#17977)
 (#19447)

* Relaxing type requirements for broadcast_like

* enhance unit test

Co-authored-by: wicky <tbc.dengwenqi@gmail.com>
---
 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 daea4b2046d9..84e7f7bd31a2 100644
--- 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
index c4dd68fcb73d..92fd030ba278 100644
--- 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):