diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 65103d786cda..4b1b2817cc7e 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -244,13 +244,14 @@ def get_info( shape = [] shape_name = [] for dim in info_proto.type.tensor_type.shape.dim: - name = dim.dim_param - value = dim.dim_value - if value is None or value == 0: - value = parse_shape_name(name, value_dict) - shape_name.append(name) - else: + if dim.HasField("dim_value"): + value = dim.dim_value shape_name.append(value) + else: + name = dim.dim_param + # Each anonymous extent is independent. Cache only named dimensions. + value = parse_shape_name(name, value_dict) if name else tirx.Var("", "int64") + shape_name.append(name) shape.append(value) name = info_proto.name diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 7f540a53f6f9..88dca16ffca7 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -4118,6 +4118,46 @@ def collect_ops(expr): assert "relax.shape_of" in op_names +@pytest.mark.skipif(not tvm_tirx.testing.device_enabled("llvm"), reason="LLVM is required") +@pytest.mark.parametrize("first_dim", [None, 0]) +def test_independent_anonymous_input_dimensions(first_dim): + x_dims = [first_dim, None, "batch"] + z_dims = ["batch", None] + graph = helper.make_graph( + [helper.make_node("Identity", ["x"], ["y"]), helper.make_node("Identity", ["z"], ["w"])], + "independent_anonymous_input_dimensions", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, x_dims), + helper.make_tensor_value_info("z", TensorProto.FLOAT, z_dims), + ], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, x_dims), + helper.make_tensor_value_info("w", TensorProto.FLOAT, z_dims), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + onnx.checker.check_model(model, full_check=True) + mod = from_onnx(model) + + # Different anonymous extents are legal within and across inputs. + extent = 2 if first_dim is None else first_dim + x = np.arange(extent * 3 * 4, dtype=np.float32).reshape(extent, 3, 4) + z = np.arange(20, dtype=np.float32).reshape(4, 5) + executable = tvm.compile(mod, target="llvm") + vm = relax.VirtualMachine(executable, tvm_tirx.cpu()) + y, w = vm["main"](tvm_tirx.runtime.tensor(x), tvm_tirx.runtime.tensor(z)) + np.testing.assert_array_equal(y.numpy(), x, strict=True) + np.testing.assert_array_equal(w.numpy(), z, strict=True) + + x_shape, z_shape = [param.ty.shape.values for param in mod["main"].params] + assert x_shape[2].same_as(z_shape[0]) + assert not x_shape[1].same_as(z_shape[1]) + if first_dim is None: + assert not x_shape[0].same_as(x_shape[1]) + else: + assert isinstance(x_shape[0], tvm_tirx.ir.IntImm) and x_shape[0].value == 0 + + def test_slice_of_scalar_shape(): # Slice consuming Shape of a rank-0 input used to raise "Slice requires a # statically known input rank", because Shape handed it an opaque value