This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 0eaf1cb019 [Relax][ONNX] Import Min/Max/Sum/Mean when an input has no
static shape (#20288)
0eaf1cb019 is described below
commit 0eaf1cb019f8bd2bcb225b92aaa36d2878b19a16
Author: Arpit Jain <[email protected]>
AuthorDate: Tue Sep 8 00:39:06 2026 -0400
[Relax][ONNX] Import Min/Max/Sum/Mean when an input has no static shape
(#20288)
Fixes #20280.
`MultiInputBase._impl_v1` reads `inp.ty.shape` for every input and folds
them with `compute_broadcast_shape`, which starts with `len(shape_a)`.
Relax spells an unknown static shape `R.Tensor(dtype=..., ndim=k)`,
whose struct info carries `shape=None`, so `len(None)` raises
`TypeError: object of type 'NoneType' has no len()`. That state is not
exotic: `R.dynamic_strided_slice` produces it, so a plain ONNX `Slice`
with runtime `starts`/`ends` feeding a `Max` is enough, on a model
`onnx.checker.check_model(..., full_check=True)` accepts and onnxruntime
executes.
ONNX defines Min, Max, Sum and Mean as elementwise with multidirectional
broadcasting, so when a static target shape cannot be computed the ops
are folded pairwise instead and the binary op broadcasts. `Mean` divides
by the input count afterwards. The known-shape path is untouched, and
single-input nodes now come back as the input itself rather than
tripping the `Found null pointer node` assertion that the
stack-and-reduce path hits with an unknown shape.
Verification. I do not have a source build handy, so I ran the release
wheel (`apache-tvm` 0.26.0, arm64) with this file's `main` version of
`_normalize_shape_dim`, `_broadcast_shape_dims`,
`compute_broadcast_shape` and `MultiInputBase` transplanted in, once
without the change and once with it. That runs `main`'s exact Python
logic for this path against a real runtime.
Across `{Min, Max, Sum, Mean}` x `{1, 2, 3}` inputs, all twelve fail
before and all twelve import after:
Max n=1: FAILED InternalError: Check failed: (n.defined()) is false:
Found null pointer node
Max n=2: FAILED TypeError: object of type 'NoneType' has no len()
The new `test_multi_input_unknown_static_shape` covers the same matrix
and shows 12 failed / 12 passed the same way. The existing `-k
multi_input` tests are 37 passed with the change in place, so the static
and symbolic paths are unaffected.
One thing worth being explicit about: this fixes the import, not
end-to-end lowering. The imported IR is `R.maximum(lv8, y)` where `lv8`
is `R.Tensor(dtype="float32", ndim=2)`, and `tvm.compile` on it still
fails with `CodeGenVM cannot handle this intrinsic now: relax.maximum`.
That is not new or specific to these ops: the same graph with `Add` in
place of `Max` imports today and fails at codegen identically with
`relax.add`. Lowering elementwise ops over an unknown-shape operand is a
separate gap; this change gets Min/Max/Sum/Mean to the same place `Add`
already is, which is what the issue asks for.
Signed-off-by: Arpit Jain <[email protected]>
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 37 ++++++++++++++++++++++++
tests/python/relax/test_frontend_onnx.py | 38 +++++++++++++++++++++++++
2 files changed, 75 insertions(+)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 2a9cea943c..861dc840f5 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -2697,6 +2697,25 @@ class MultiInputBase(OnnxOpConverter):
numpy_op: Callable = None
relax_op: Callable = None
+ # Pairwise equivalent, used when no static broadcast shape can be computed.
+ binary_op: Callable = None
+
+ @classmethod
+ def _impl_dynamic(cls, bb, inputs):
+ """Fold the inputs pairwise, letting the binary op broadcast.
+
+ ONNX defines Min, Max, Sum and Mean as elementwise with
multidirectional
+ broadcasting, so the pairwise form is equivalent to the
stack-and-reduce
+ form and does not need a shape known at import time.
+ """
+ if cls.binary_op is None:
+ raise NotImplementedError(
+ f"{cls.__name__} cannot import an input whose static shape is
unknown"
+ )
+ return functools.reduce(
+ lambda lhs, rhs: bb.normalize(cls.binary_op(lhs, rhs)), # pylint:
disable=not-callable
+ inputs,
+ )
@classmethod
def _impl_v1(cls, bb, inputs, attr, params):
@@ -2718,6 +2737,15 @@ class MultiInputBase(OnnxOpConverter):
return relax.const(output, output.dtype)
input_shapes = [inp.ty.shape for inp in inputs]
+ if any(shape is None for shape in input_shapes):
+ # Relax spells an unknown static shape R.Tensor(dtype=..., ndim=k),
+ # whose struct info carries shape None. R.dynamic_strided_slice
+ # produces exactly that, so a plain ONNX Slice with runtime
+ # starts/ends reaches here and compute_broadcast_shape raised
+ # `object of type 'NoneType' has no len()` on a model onnx.checker
+ # accepts and onnxruntime runs.
+ return cls._impl_dynamic(bb, inputs)
+
target_shape = functools.reduce(compute_broadcast_shape, input_shapes)
# broadcast_to, stack them, then perform minimum over the new axis.
@@ -2731,6 +2759,7 @@ class Min(MultiInputBase):
numpy_op = _np.min
relax_op = relax.op.min
+ binary_op = relax.op.minimum
class Max(MultiInputBase):
@@ -2738,6 +2767,7 @@ class Max(MultiInputBase):
numpy_op = _np.max
relax_op = relax.op.max
+ binary_op = relax.op.maximum
class Mean(MultiInputBase):
@@ -2745,6 +2775,12 @@ class Mean(MultiInputBase):
numpy_op = _np.mean
relax_op = relax.op.mean
+ binary_op = relax.op.add
+
+ @classmethod
+ def _impl_dynamic(cls, bb, inputs):
+ total = super()._impl_dynamic(bb, inputs)
+ return relax.op.divide(total, relax.const(len(inputs),
inputs[0].ty.dtype))
class Sum(MultiInputBase):
@@ -2752,6 +2788,7 @@ class Sum(MultiInputBase):
numpy_op = _np.sum
relax_op = relax.op.sum
+ binary_op = relax.op.add
class Log(OnnxOpConverter):
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 94268d477e..335694ed76 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -1216,6 +1216,44 @@ def test_multi_input_constant_single_input(op_name,
shape):
check_correctness(helper.make_model(graph), opset=13)
[email protected]("op_name", ["Min", "Max", "Sum", "Mean"])
[email protected]("num_inputs", [1, 2, 3])
+def test_multi_input_unknown_static_shape(op_name, num_inputs):
+ """An input with no static shape imports instead of raising len(None).
+
+ A Slice with runtime starts/ends lowers to R.dynamic_strided_slice, whose
+ struct info is R.Tensor(dtype=..., ndim=k) with shape None. That reaches
+ MultiInputBase, where compute_broadcast_shape used to call len() on it.
+ The model is valid ONNX and onnxruntime executes it.
+ """
+ slice_node = helper.make_node(
+ "Slice", ["x", "starts", "ends", "axes"], ["sliced"], name="slice0"
+ )
+ other_names = [f"y{i}" for i in range(num_inputs - 1)]
+ op_node = helper.make_node(op_name, ["sliced"] + other_names, ["output"],
name="op0")
+
+ graph = helper.make_graph(
+ [slice_node, op_node],
+ f"slice_then_{op_name.lower()}",
+ inputs=[
+ helper.make_tensor_value_info("x", TensorProto.FLOAT, [4, 3]),
+ helper.make_tensor_value_info("starts", TensorProto.INT64, [1]),
+ helper.make_tensor_value_info("ends", TensorProto.INT64, [1]),
+ helper.make_tensor_value_info("axes", TensorProto.INT64, [1]),
+ ]
+ + [
+ helper.make_tensor_value_info(name, TensorProto.FLOAT, [2, 3])
+ for name in other_names
+ ],
+ outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT,
[2, 3])],
+ )
+ model = helper.make_model(graph, opset_imports=[helper.make_opsetid("",
13)])
+ onnx.checker.check_model(model, full_check=True)
+
+ tvm_model = from_onnx(model, keep_params_in_input=True)
+ assert "dynamic_strided_slice" in str(tvm_model)
+
+
@pytest.mark.parametrize("op_name", ["And", "Or", "Xor"])
def test_binary_bool(op_name: str):
verify_binary(op_name, [32, 32], [32, 32], [32, 32],
dtype=TensorProto.BOOL)
