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{"lhs", "rhs"}; }) -.set_attr("FInferType", ElemwiseType<2, 1>) +.set_attr("FInferType", [](const nnvm::NodeAttrs& attrs, + std::vector *in_attrs, + std::vector *out_attrs) { + CHECK_EQ(in_attrs->size(), 2) << " in operator " << attrs.name; + std::vector 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("FGradient", [](const nnvm::ObjectPtr& n, const std::vector& 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):