[Bug][Frontend][ONNX] Failure of test_forward.test_aten unit test after upgrade to PyTorch 1.12.0

[oviazlo]

As described in the discussion forum thread the test below fails after the upgrade to the PyTorch 1.12.0:

tvm/tests/python/frontend/onnx/test_forward.py

Lines 5471 to 5507 in 458ca81

	@tvm.testing.parametrize_targets
	def test_aten(target, dev):
	"""test_aten"""
	torch.set_grad_enabled(False)
	def _convert_to_onnx(model, inputs):
	file_name = "aten_model.onnx"
	torch.onnx.export(
	model,
	inputs,
	file_name,
	export_params=True,
	verbose=False,
	opset_version=10,
	operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN,
	)
	onnx_model = onnx.load(file_name)
	assert 's: "embedding_bag"' in str(onnx_model)
	return onnx_model
	def verify_embedding_bag(num_embedding, embedding_dim, data_shape, num_bags=None):
	dummy_data = torch.randint(0, num_embedding - 1, data_shape)
	tvm_inputs = [dummy_data.numpy()]
	model = torch.nn.EmbeddingBag(num_embedding, embedding_dim)
	onnx_model = _convert_to_onnx(model, dummy_data)
	torch_out = model(dummy_data)
	tvm_out = get_tvm_output_with_vm(
	onnx_model,
	tvm_inputs,
	freeze_params=True,
	target=target,
	dev=dev,
	)
	tvm.testing.assert_allclose(torch_out.numpy(), tvm_out, atol=5e-7)
	verify_embedding_bag(10, 3, [2, 10])
	verify_embedding_bag(32, 2, [3, 3])

The test fails in the line:

tvm/tests/python/frontend/onnx/test_forward.py

Line 5488 in 458ca81

assert 's: "embedding_bag"' in str(onnx_model)

The problem lies due to improper model conversion from PyTorch to ONNX with PyTorch 1.12.0. It looks like that argument below is ignored:

tvm/tests/python/frontend/onnx/test_forward.py

Line 5485 in 458ca81

operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN,

and ATen PyTorch operators are substituted with native ONNX operators (instead of being exported as ATen ops in ONNX). This behaviour fail an assert above.

I have created an appropriate bug report in the PyTorch repo: pytorch/pytorch#87313

One possible quick solution will be to remove an assert above in _convert_to_onnx function (since the test produces correct results and satisfies all other asserts). However, it will reduce coverage of the test because onnx_model (in the code below) will not contain ATen operators, and the possibility to consume ONNX graph with such operators by TVM will not be tested.

tvm/tests/python/frontend/onnx/test_forward.py

Lines 5497 to 5503 in 458ca81

	tvm_out = get_tvm_output_with_vm(
	onnx_model,
	tvm_inputs,
	freeze_params=True,
	target=target,
	dev=dev,
	)

Environment

• TVM version: v0.10.0rc0 (82c4add)
• OS: macOS 12.6 (arm64)
• torch 1.12.0+cpu
• torchvision 0.13.0+cpu

Steps to reproduce

With Torch 1.12.0 and Torchvision 0.13.0 environment, run:
python -m pytest tests/python/frontend/onnx/test_forward.py

Triage

• frontend:onnx

[oviazlo]

@masahi FYI

[masahi]

Since "ATen" op in an ONNX graph is never meant to be an official thing, I think it is fine to remove the broken assert. Or disable the test entirely for PT 1.12 and newer.