已合并
fix issues with npu_format_cast function parameters #32693
Kuteriod创建于 3月30日
fix issues with npu_format_cast function parameters #32693
已合并
Kuteriod创建于 3月30日
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- @unittest.skip
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- @unittest.skip
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.skip907+ 
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- @unittest.skip
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- @unittest.skip
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):
258class _NPUFormatCastOP(torch.autograd.Function):258class _NPUFormatCastOP(torch.autograd.Function):
259 259 
260 @staticmethod260 @staticmethod
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 @staticmethod265 @staticmethod
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