diff --git a/python/tvm/relay/frontend/mxnet.py b/python/tvm/relay/frontend/mxnet.py index 712c025fd1b6..05adef3dece3 100644 --- a/python/tvm/relay/frontend/mxnet.py +++ b/python/tvm/relay/frontend/mxnet.py @@ -2230,6 +2230,16 @@ def _mx_broadcast_to(inputs, attrs): return _op.broadcast_to(data, tgt_shape) +def _mx_broadcast_like(inputs, attrs): + assert len(inputs) == 2 + for axes in ["lhs_axes", "rhs_axes"]: + if axes in attrs.attrs: + raise tvm.error.OpAttributeUnImplemented( + 'Attribute "{}" is not supported for operator broadcast_like.'.format(axes) + ) + return _op.broadcast_to_like(*inputs) + + def _mx_logical_not(inputs, input_types): data = inputs[0] dtype = _infer_type(data).checked_type.dtype @@ -2389,6 +2399,7 @@ def _mx_npi_where_rscalar(inputs, attrs): "broadcast_logical_and": _mx_broadcast_logical(_op.logical_and), "broadcast_logical_xor": _mx_broadcast_logical(_op.logical_xor), "broadcast_to": _mx_broadcast_to, + "broadcast_like": _mx_broadcast_like, "logical_not": _mx_logical_not, "_equal": _mx_compare(_op.equal, _rename), "_not_equal": _mx_compare(_op.not_equal, _rename), diff --git a/tests/python/frontend/mxnet/test_forward.py b/tests/python/frontend/mxnet/test_forward.py index 52c5db8b170f..5c9d1ca8d8a5 100644 --- a/tests/python/frontend/mxnet/test_forward.py +++ b/tests/python/frontend/mxnet/test_forward.py @@ -715,6 +715,24 @@ def verify(input_shape, shape): verify((4, 1, 32, 32), (4, 8, 32, 32)) +@tvm.testing.uses_gpu +def test_forward_broadcast_like(): + def verify(input_shape, like_shape): + x_np = np.random.uniform(size=input_shape).astype("float32") + y_np = np.random.uniform(size=like_shape).astype("float32") + ref_res = mx.nd.broadcast_like(mx.nd.array(x_np), mx.nd.array(y_np)) + mx_sym = mx.sym.broadcast_like(mx.sym.var("x"), mx.sym.var("y")) + mod, _ = relay.frontend.from_mxnet(mx_sym, {"x": input_shape, "y": like_shape}) + for target, ctx in tvm.testing.enabled_targets(): + for kind in ["graph", "debug"]: + intrp = relay.create_executor(kind, mod=mod, ctx=ctx, target=target) + op_res = intrp.evaluate()(x_np, y_np) + tvm.testing.assert_allclose(op_res.asnumpy(), ref_res.asnumpy()) + + verify((1, 2, 3), (3, 2, 3)) + verify((4, 1, 32, 32), (4, 8, 32, 32)) + + @tvm.testing.uses_gpu def test_forward_logical_not(): a_shape = (3, 4, 5)