[BUGFIX] Fix broadcast_like

src/api/operator/numpy_extension/npx_broadcast_like_op.cc

@@ -44,14 +44,18 @@ MXNET_REGISTER_API("_npx.broadcast_like")
   // lhs_axes
   if (args[2].type_code() == kNull) {
     param.lhs_axes = dmlc::optional<mxnet::TShape>();
+  } else if (args[2].type_code() == kDLInt) {
+    param.lhs_axes = TShape(1, args[2].operator int64_t());
   } else {
     param.lhs_axes = mxnet::TShape(args[2].operator ObjectRef());
   }
   // rhs_axes
-  if (args[2].type_code() == kNull) {
+  if (args[3].type_code() == kNull) {
     param.rhs_axes = dmlc::optional<mxnet::TShape>();
+  } else if (args[3].type_code() == kDLInt) {
+    param.rhs_axes = TShape(1, args[3].operator int64_t());
   } else {
-    param.rhs_axes = mxnet::TShape(args[2].operator ObjectRef());
+    param.rhs_axes = mxnet::TShape(args[3].operator ObjectRef());
   }
 
   attrs.op = op;

tests/python/unittest/test_numpy_op.py

@@ -10338,3 +10338,14 @@ def test_modulated_deformable_convolution(num_batch, num_channel_data, num_defor
         rtol, atol = 1.0, 1e-2
     else:
         rtol, atol = 0.05, 1e-3
+
+
+@use_np
+def test_broadcast_like_different_types():
+    x = mx.np.zeros((2, 1))
+    y = mx.np.ones((2, 2))
+
+    y = mx.np.array(y).astype('int32')
+    z = mx.npx.broadcast_like(x, y, 1, 1)
+    assert_almost_equal(z.asnumpy(), np.array([[0,0],[0,0]]))
+    assert x.dtype == z.dtype