已合并
fix issues with npu_format_cast function parameters #32693
Kuteriod创建于 3月30日
fix issues with npu_format_cast function parameters #32693
已合并
共 3 个文件变更+6-9
| @@ -203,7 +203,6 @@ class TestOnnxOps(TestCase): | |||
| 203 | assert(os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, | 203 | assert(os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, |
| 204 | onnx_model_name))) | 204 | onnx_model_name))) |
| 205 | 205 | ||
| 206 | - | ||
| 207 | def test_wrapper_npu_multi_head_attention(self): | 206 | def test_wrapper_npu_multi_head_attention(self): |
| 208 | class Model(torch.nn.Module): | 207 | class Model(torch.nn.Module): |
| 209 | def __init__(self): | 208 | def __init__(self): |
| @@ -387,7 +386,6 @@ class TestOnnxOps(TestCase): | |||
| 387 | assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, | 386 | assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, |
| 388 | onnx_model_name))) | 387 | onnx_model_name))) |
| 389 | 388 | ||
| 390 | - | ||
| 391 | def test_wrapper_npu_format_cast(self): | 389 | def test_wrapper_npu_format_cast(self): |
| 392 | class Model(torch.nn.Module): | 390 | class Model(torch.nn.Module): |
| 393 | def __init__(self): | 391 | def __init__(self): |
| @@ -906,7 +904,7 @@ class TestOnnxOps(TestCase): | |||
| 906 | export_onnx(onnx_model_name) | 904 | export_onnx(onnx_model_name) |
| 907 | assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, | 905 | assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, |
| 908 | onnx_model_name))) | 906 | onnx_model_name))) |
| 909 | - @unittest.skip | 907 | + |
| 910 | def test_wrapper_npu_lstm_cell(self): | 908 | def test_wrapper_npu_lstm_cell(self): |
| 911 | class Model(torch.nn.Module): | 909 | class Model(torch.nn.Module): |
| 912 | def __init__(self): | 910 | def __init__(self): |
| @@ -958,7 +956,6 @@ class TestOnnxOps(TestCase): | |||
| 958 | assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, | 956 | assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, |
| 959 | onnx_model_name))) | 957 | onnx_model_name))) |
| 960 | 958 | ||
| 961 | - | ||
| 962 | def test_wrapper_npu_lstm(self): | 959 | def test_wrapper_npu_lstm(self): |
| 963 | class Model(torch.nn.Module): | 960 | class Model(torch.nn.Module): |
| 964 | def __init__(self): | 961 | def __init__(self): |
| @@ -1015,7 +1012,6 @@ class TestOnnxOps(TestCase): | |||
| 1015 | assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, | 1012 | assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, |
| 1016 | onnx_model_name))) | 1013 | onnx_model_name))) |
| 1017 | 1014 | ||
| 1018 | - | ||
| 1019 | def test_wrapper_npu_gru(self): | 1015 | def test_wrapper_npu_gru(self): |
| 1020 | class Model(torch.nn.Module): | 1016 | class Model(torch.nn.Module): |
| 1021 | def __init__(self): | 1017 | def __init__(self): |
| @@ -2130,7 +2130,7 @@ | |||
| 2130 | "signature": "(self, other, offset1, offset2, c1_len)" | 2130 | "signature": "(self, other, offset1, offset2, c1_len)" |
| 2131 | }, | 2131 | }, |
| 2132 | "torch_npu.npu_format_cast": { | 2132 | "torch_npu.npu_format_cast": { |
| 2133 | - "signature": "(self, acl_format, customize_dtype=None)" | 2133 | + "signature": "(self, acl_format, *, customize_dtype=None)" |
| 2134 | }, | 2134 | }, |
| 2135 | "torch_npu.npu_fused_attention_score": { | 2135 | "torch_npu.npu_fused_attention_score": { |
| 2136 | "signature": "(query_layer, key_layer, value_layer, attention_mask, scale, keep_prob, query_transpose=False, key_transpose=False, bmm_score_transpose_a=False, bmm_score_transpose_b=False, value_transpose=False, dx_transpose=False)" | 2136 | "signature": "(query_layer, key_layer, value_layer, attention_mask, scale, keep_prob, query_transpose=False, key_transpose=False, bmm_score_transpose_a=False, bmm_score_transpose_b=False, value_transpose=False, dx_transpose=False)" |
| @@ -258,8 +258,9 @@ class _NPUDeformableConv2dOP(torch.autograd.Function): | |||
| 258 | class _NPUFormatCastOP(torch.autograd.Function): | 258 | class _NPUFormatCastOP(torch.autograd.Function): |
| 259 | 259 | ||
| 260 | 260 | ||
| 261 | - def forward(ctx, *args, **kwargs): | 261 | + def forward(ctx, self, acl_format, customize_dtype=None): |
| 262 | - return torch.ops.npu.npu_format_cast(*args, **kwargs) | 262 | + return torch.ops.npu.npu_format_cast( |
| 263 | + self, acl_format, customize_dtype=customize_dtype) | ||
| 263 | 264 | ||
| 264 | 265 | ||
| 265 | def symbolic(g, self: Tensor, acl_format: int, customize_dtype: int = None): | 266 | def symbolic(g, self: Tensor, acl_format: int, customize_dtype: int = None): |
| @@ -1065,7 +1066,7 @@ def _wrapper_npu_deformable_conv2d(inputs, weight, offset, bias, kernel_size, st | |||
| 1065 | padding, dilation, groups, deformable_groups, modulated) | 1066 | padding, dilation, groups, deformable_groups, modulated) |
| 1066 | 1067 | ||
| 1067 | 1068 | ||
| 1068 | -def _wrapper_npu_format_cast(self, acl_format, customize_dtype=None): | 1069 | +def _wrapper_npu_format_cast(self, acl_format, *, customize_dtype=None): |
| 1069 | return _NPUFormatCastOP.apply(self, acl_format, customize_dtype) | 1070 | return _NPUFormatCastOP.apply(self, acl_format, customize_dtype) |
| 1070 | 1071 | ||
| 1071 | 1072 | ||