From cc6ed203954d58d189e214b556868b319b8449f9 Mon Sep 17 00:00:00 2001 From: FFChopon <3483776996@qq.com> Date: Tue, 18 Aug 2026 02:56:29 +0800 Subject: [PATCH] [Relax][Frontend][ONNX] Fix Mean/Sum/Min/Max with all-constant inputs The MultiInputBase constant-fold path called numpy_op(*np_inputs) (np.mean/np.sum/np.min/np.max). For these numpy reductions the 2nd and later positional arguments are interpreted as `axis`, not as additional data tensors, so a single constant input returned a 0-d scalar (the global mean instead of the input unchanged) and multiple constant inputs raised TypeError. Fix by mirroring the non-constant path: broadcast each constant to the common shape, stack along a new leading axis, then reduce along it. Co-Authored-By: Claude --- python/tvm/relax/frontend/onnx/onnx_frontend.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 6bbf220dfe2e..ce4a9850e225 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -2455,7 +2455,20 @@ def _impl_v1(cls, bb, inputs, attr, params): raise NotImplementedError("numpy_op and relax_op must be defined for MultiInputBase") if all([isinstance(inp, relax.Constant) for inp in inputs]): np_inputs = [inp.data.numpy() for inp in inputs] - output = cls.numpy_op(*np_inputs) # pylint: disable=not-callable + # numpy_op (np.mean/np.sum/np.min/np.max) reduces its first arg, + # treating any further positional args as `axis`, so calling it as + # numpy_op(*np_inputs) is wrong for the variadic ONNX semantics. + # Broadcast to the common shape, stack along a new leading axis, + # then reduce along it — mirrors the non-constant path below. + input_shapes = [inp.ty.shape for inp in inputs] + target_shape = tuple( + int(dim) + for dim in functools.reduce(compute_broadcast_shape, input_shapes) + ) + stacked = _np.stack( + [_np.broadcast_to(x, target_shape) for x in np_inputs], axis=0 + ) + output = cls.numpy_op(stacked, axis=0) # pylint: disable=not-callable return relax.const(output, output.dtype) input_shapes = [inp.ty.shape for inp in inputs]