[Relax][PyTorch] Fix output shape of torch.nn.functional.scaled_dot_product_attention

[mshr-h]

torch.nn.functional.scaled_dot_product_attention outputs in the shape of (N, ..., L, E_v) but relax.op.nn.attention does (N, L, ..., E_v) so the output should also be transposed.

Maybe we should add E2E tests in tests/python/nightly/ to check the relax torch frontend.

cc: @yongwww

[yongwww]

we can transpose to get the expected result. Thanks for the effort!

[yongwww]

overall looks good to me

[mshr-h]

MSC E2E test is failing. Seems like we also need to change something other than relax frontend.
@Archermmt Do you have any ideas on how to fix the error?

Link to the ci log: https://ci.tlcpack.ai/blue/organizations/jenkins/tvm-unity/detail/PR-17379/6/pipeline/

tests/python/contrib/test_msc/test_translate_torch.py::test_attention FAILED

[2024-09-17T10:06:53.212Z] tests/python/contrib/test_msc/test_translate_torch.py::test_attention FAILED
[2024-09-17T10:06:53.212Z]
[2024-09-17T10:06:53.212Z] =================================== FAILURES ===================================
[2024-09-17T10:06:53.212Z] ________________________________ test_attention ________________________________
[2024-09-17T10:06:53.212Z]
[2024-09-17T10:06:53.212Z] def test_attention():
[2024-09-17T10:06:53.212Z] """test torch translator for attention"""
[2024-09-17T10:06:53.212Z]
[2024-09-17T10:06:53.212Z] # pylint: disable=import-outside-toplevel
[2024-09-17T10:06:53.212Z] import torch.nn.functional as F
[2024-09-17T10:06:53.212Z]
[2024-09-17T10:06:53.212Z] class Attention1(Module):
[2024-09-17T10:06:53.212Z] def forward(self, q_data, k_data, v_data):
[2024-09-17T10:06:53.212Z] return F.scaled_dot_product_attention(q_data, k_data, v_data)
[2024-09-17T10:06:53.212Z]
[2024-09-17T10:06:53.212Z] class Attention2(Module):
[2024-09-17T10:06:53.212Z] def forward(self, q_data, k_data, v_data):
[2024-09-17T10:06:53.212Z] return F.scaled_dot_product_attention(q_data, k_data, v_data, is_causal=True)
[2024-09-17T10:06:53.212Z]
[2024-09-17T10:06:53.212Z] input_info = [
[2024-09-17T10:06:53.212Z] ([32, 8, 128, 64], "float32"),
[2024-09-17T10:06:53.212Z] ([32, 8, 128, 64], "float32"),
[2024-09-17T10:06:53.212Z] ([32, 8, 128, 64], "float32"),
[2024-09-17T10:06:53.212Z] ]
[2024-09-17T10:06:53.212Z] > verify_model(Attention1(), input_info)
[2024-09-17T10:06:53.212Z]
[2024-09-17T10:06:53.212Z] tests/python/contrib/test_msc/test_translate_torch.py:1127:
[2024-09-17T10:06:53.212Z] _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
[2024-09-17T10:06:53.212Z] tests/python/contrib/test_msc/test_translate_torch.py:52: in verify_model
[2024-09-17T10:06:53.212Z] tvm.testing.assert_allclose(
[2024-09-17T10:06:53.212Z] _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
[2024-09-17T10:06:53.212Z]
[2024-09-17T10:06:53.212Z] actual = array([[[[0.5253085 , 0.48211107, 0.52921075, ..., 0.518867 ,
[2024-09-17T10:06:53.212Z] 0.49926636, 0.48493868],
[2024-09-17T10:06:53.212Z] [0.5294311 ...5],
[2024-09-17T10:06:53.212Z] [0.47335747, 0.48579183, 0.5360674 , ..., 0.543607 ,
[2024-09-17T10:06:53.212Z] 0.5020893 , 0.47848547]]]], dtype=float32)
[2024-09-17T10:06:53.212Z] desired = array([[[[0.5253085 , 0.48211107, 0.52921075, ..., 0.518867 ,
[2024-09-17T10:06:53.212Z] 0.49926636, 0.48493868],
[2024-09-17T10:06:53.212Z] [0.49697113... ],
[2024-09-17T10:06:53.212Z] [0.47335747, 0.48579183, 0.5360674 , ..., 0.543607 ,
[2024-09-17T10:06:53.212Z] 0.5020893 , 0.47848547]]]], dtype=float32)
[2024-09-17T10:06:53.212Z] rtol = 1e-05, atol = 1e-05
[2024-09-17T10:06:53.212Z]
[2024-09-17T10:06:53.212Z] def assert_allclose(actual, desired, rtol=1e-7, atol=1e-7):
[2024-09-17T10:06:53.212Z] """Version of np.testing.assert_allclose with atol and rtol fields set
[2024-09-17T10:06:53.212Z] in reasonable defaults.
[2024-09-17T10:06:53.212Z]
[2024-09-17T10:06:53.212Z] Arguments actual and desired are not interchangeable, since the function
[2024-09-17T10:06:53.212Z] compares the abs(actual-desired) with atol+rtol*abs(desired). Since we
[2024-09-17T10:06:53.212Z] often allow desired to be close to zero, we generally want non-zero atol.
[2024-09-17T10:06:53.212Z] """
[2024-09-17T10:06:53.212Z] actual = np.asanyarray(actual)
[2024-09-17T10:06:53.212Z] desired = np.asanyarray(desired)
[2024-09-17T10:06:53.212Z] > np.testing.assert_allclose(actual.shape, desired.shape)
[2024-09-17T10:06:53.212Z] E AssertionError:
[2024-09-17T10:06:53.212Z] E Not equal to tolerance rtol=1e-07, atol=0
[2024-09-17T10:06:53.212Z] E
[2024-09-17T10:06:53.212Z] E Mismatched elements: 2 / 4 (50%)
[2024-09-17T10:06:53.212Z] E Max absolute difference: 120
[2024-09-17T10:06:53.212Z] E Max relative difference: 15.
[2024-09-17T10:06:53.212Z] E x: array([ 32, 8, 128, 64])
[2024-09-17T10:06:53.212Z] E y: array([ 32, 128, 8, 64])
[2024-09-17T10:06:53.212Z]
[2024-09-17T10:06:53.212Z] python/tvm/testing/utils.py:119: AssertionError