mindstudio-probe 1.0.4__py3-none-any.whl → 1.1.1__py3-none-any.whl

msprobe/pytorch/free_benchmark/result_handlers/fix_handler.py

@@ -1,9 +1,25 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Any
 
+from msprobe.core.common.exceptions import FreeBenchmarkException
+from msprobe.pytorch.free_benchmark import logger
 from msprobe.pytorch.free_benchmark.common.params import DataParams
 from msprobe.pytorch.free_benchmark.common.utils import Tools
 from msprobe.pytorch.free_benchmark.result_handlers.base_handler import FuzzHandler
-from msprobe.pytorch.free_benchmark import logger
 
 
 class FixHandler(FuzzHandler):
@@ -16,9 +32,9 @@ class FixHandler(FuzzHandler):
             return Tools.convert_fuzz_output_to_origin(
                 data_params.original_result, data_params.perturbed_result
             )
-        except Exception as e:
-            logger.warning_on_rank_0(
+        except FreeBenchmarkException as e:
+            logger.warning(
                 f"[msprobe] Free Benchmark: For {self.params.api_name} "
-                f"Fix output failed. "
+                f"Fix output failed because of: \n{e}"
             )
-        return data_params.original_result
+            return data_params.original_result

msprobe/pytorch/free_benchmark/result_handlers/handler_factory.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from msprobe.pytorch.free_benchmark import FreeBenchmarkException
 from msprobe.pytorch.free_benchmark.common.constant import PreheatConfig
 from msprobe.pytorch.free_benchmark.common.enums import HandlerType

msprobe/pytorch/free_benchmark/result_handlers/preheat_handler.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import math
 from typing import Any
 
@@ -118,8 +133,10 @@ class PreheatHandler(FuzzHandler):
         """
         # 每一步样本数
         total_count = preheat_counter.get_one_step_used_api(self.pure_name)
-        sample_count_per_step = self._get_sample_count_per_step()
         need_sample_set = set()
+        if total_count == 0:
+            return need_sample_set
+        sample_count_per_step = self._get_sample_count_per_step()
         prehead_step = self.params.preheat_config.get("preheat_step")
         for i in range(1, sample_count_per_step + 1):
             count = (prehead_step * (i - 1) + self.params.step) % total_count
@@ -136,9 +153,7 @@ class PreheatHandler(FuzzHandler):
 
     def _adjust_threshold_for_dtype(self, dtype_str, compare_result):
         con_ratio = [ratio for ratio, is_consistent in compare_result if is_consistent]
-        incon_ratio = [
-            ratio for ratio, is_consistent in compare_result if not is_consistent
-        ]
+        incon_ratio = [ratio for ratio, is_consistent in compare_result if not is_consistent]
         old_thd = preheat_counter.get_api_thd(self.pure_name, dtype_str)
         new_thd = old_thd
         # 正例负例都存在

msprobe/pytorch/function_factory.py

@@ -1,4 +1,18 @@
-from msprobe.pytorch.common.utils import logger
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from msprobe.pytorch.bench_functions.apply_adam_w import npu_apply_adam_w
 from msprobe.pytorch.bench_functions.confusion_transpose import npu_confusion_transpose, \
     npu_confusion_transpose_backward
@@ -12,7 +26,8 @@ from msprobe.pytorch.bench_functions.rms_norm import npu_rms_norm, npu_rms_norm_
 from msprobe.pytorch.bench_functions.rotary_mul import npu_rotary_mul, npu_rotary_mul_backward
 from msprobe.pytorch.bench_functions.scaled_mask_softmax import npu_scaled_masked_softmax, \
     npu_scaled_masked_softmax_backward
-from msprobe.pytorch.bench_functions.swiglu import npu_swiglu, npu_swiglu_backward, swish_grad, swish
+from msprobe.pytorch.bench_functions.swiglu import npu_swiglu, npu_swiglu_backward
+from msprobe.pytorch.common.utils import logger
 
 
 class Register(dict):

msprobe/pytorch/functional/module_dump.py

@@ -0,0 +1,84 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+from msprobe.core.common.const import Const
+from msprobe.core.common.exceptions import MsprobeException
+from msprobe.core.data_dump.scope import BaseScope
+from msprobe.pytorch.common.log import logger
+from msprobe.pytorch.debugger.precision_debugger import PrecisionDebugger
+from msprobe.pytorch.hook_module.api_registry import api_register
+from msprobe.pytorch.service import torch_version_above_or_equal_2
+
+hook_handle_list = []
+
+
+def module_dump(module, dump_name):
+    if not isinstance(module, nn.Module):
+        logger.error("The parameter module in module_dump must be a Module subclass.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+    if not isinstance(dump_name, str):
+        logger.error("The parameter dump_name in module_dump must be a str type.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+
+    api_register.api_originality()
+    register_hook(module, dump_name)
+
+
+def module_dump_end():
+    api_register.api_modularity()
+    remove_hook()
+    hook_handle_list.clear()
+
+
+def register_hook(module, dump_name):
+    prefix = BaseScope.Module_Type_Module + Const.SEP + dump_name + Const.SEP + module.__class__.__name__ + Const.SEP
+
+    pdg = PrecisionDebugger()
+    _, forward_hook, backward_hook, forward_hook_torch_version_below_2 = \
+        pdg.service.build_hook(BaseScope.Module_Type_Module, prefix)
+
+    if torch_version_above_or_equal_2:
+        forward_hook_handle = module.register_forward_hook(forward_hook, with_kwargs=True)
+        hook_handle_list.append(forward_hook_handle)
+    else:
+        pdg.service.check_register_full_backward_hook(module)
+        full_backward_hook_handle = module.register_full_backward_hook(
+            pdg.service.module_processor.node_hook(prefix + Const.BACKWARD, Const.STOP))
+        forward_hook_handle = module.register_forward_hook(forward_hook_torch_version_below_2)
+        hook_handle_list.extend([full_backward_hook_handle, forward_hook_handle])
+    pdg.service.check_register_full_backward_hook(module)
+    full_backward_hook_handle = module.register_full_backward_hook(backward_hook)
+
+    forward_pre_hook_handle = module.register_forward_pre_hook(
+        pdg.service.module_processor.node_hook(prefix + Const.FORWARD, Const.START))
+    forward_hook_handle = module.register_forward_hook(
+        pdg.service.module_processor.node_hook(prefix + Const.FORWARD, Const.STOP))
+    hook_handle_list.extend([full_backward_hook_handle, forward_pre_hook_handle, forward_hook_handle])
+
+    if torch_version_above_or_equal_2:
+        backward_pre_hook_handle = module.register_full_backward_pre_hook(
+            pdg.service.module_processor.node_hook(prefix + Const.BACKWARD, Const.START))
+        pdg.service.check_register_full_backward_hook(module)
+        full_backward_hook_handle = module.register_full_backward_hook(
+            pdg.service.module_processor.node_hook(prefix + Const.BACKWARD, Const.STOP))
+        hook_handle_list.extend([backward_pre_hook_handle, full_backward_hook_handle])
+
+
+def remove_hook():
+    for hook_handle in hook_handle_list:
+        if isinstance(hook_handle, torch.utils.hooks.RemovableHandle):
+            hook_handle.remove()

msprobe/pytorch/grad_probe/grad_monitor.py

@@ -1,15 +1,31 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import os
 from collections import defaultdict
 
 import torch
-if int(torch.__version__.split('.')[0]) >= 2:
-    from torch.optim.optimizer import register_optimizer_step_pre_hook
-from msprobe.pytorch.grad_probe.grad_stat_csv import GradStatCsv
-from msprobe.core.grad_probe.utils import check_numeral_list_ascend, data_in_list_target
+from msprobe.core.common.file_utils import remove_path, save_npy, write_csv, create_directory
 from msprobe.core.grad_probe.constant import level_adp
+from msprobe.core.grad_probe.utils import check_numeral_list_ascend, data_in_list_target
 from msprobe.pytorch.common.log import logger
-from msprobe.core.common.file_utils import remove_path, save_npy, write_csv, create_directory
 from msprobe.pytorch.common.utils import get_rank_id, print_rank_0
+from msprobe.pytorch.grad_probe.grad_stat_csv import GradStatCsv
+
+if int(torch.__version__.split('.')[0]) >= 2:
+    from torch.optim.optimizer import register_optimizer_step_pre_hook
 
 
 class GradientMonitor:
@@ -75,7 +91,7 @@ class GradientMonitor:
                 output_lines.append(grad_info)
                 if self._level_adp["have_grad_direction"]:
                     GradientMonitor.save_grad_direction(param_name, grad,
-                                                    f'{self._output_path}/rank{self._rank}/step{self._step}')
+                                                        f'{self._output_path}/rank{self._rank}/step{self._step}')
             output_dirpath = os.path.join(self._output_path, f"rank{getattr(self, '_rank')}")
             if not os.path.isdir(output_dirpath):
                 create_directory(output_dirpath)
@@ -87,5 +103,6 @@ class GradientMonitor:
             output_lines.insert(0, header_result)
             write_csv(output_lines, output_path)
             logger.info(f"write grad data to {output_path}")
+
         if int(torch.__version__.split('.')[0]) >= 2:
             register_optimizer_step_pre_hook(optimizer_pre_step_hook)

msprobe/pytorch/grad_probe/grad_stat_csv.py

@@ -1,11 +1,27 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from abc import ABC, abstractmethod
 from collections import namedtuple
 import hashlib
+from functools import wraps
 import torch
 from msprobe.core.grad_probe.constant import GradConst
 
-CSV_header_input = namedtuple("CSV_header_input", ["bounds"])
-CSV_content_input = namedtuple("CSV_content_input", ["grad", "bounds"])
+CsvHeaderInput = namedtuple("CsvHeaderInput", ["bounds"])
+CsvContentInput = namedtuple("CsvContentInput", ["grad", "bounds"])
 
 
 class GradStatCsv:
@@ -15,7 +31,7 @@ class GradStatCsv:
     def generate_csv_header(level, bounds):
         header = ["param_name"]
         for key in level["header"]:
-            csv_header_input = CSV_header_input(bounds=bounds)
+            csv_header_input = CsvHeaderInput(bounds=bounds)
             header.extend(GradStatCsv.csv[key].generate_csv_header(csv_header_input))
         return header
 
@@ -23,7 +39,7 @@ class GradStatCsv:
     def generate_csv_line(param_name, level, grad, bounds):
         line = [param_name]
         for key in level["header"]:
-            csv_content_input = CSV_content_input(grad=grad, bounds=bounds)
+            csv_content_input = CsvContentInput(grad=grad, bounds=bounds)
             line.extend(GradStatCsv.csv[key].generate_csv_content(csv_content_input))
         return line
 
@@ -37,20 +53,24 @@ def register_csv_item(key, cls=None):
 
 
 class CsvItem(ABC):
+    @staticmethod
     @abstractmethod
     def generate_csv_header(csv_header_input):
         pass
 
+    @staticmethod
     @abstractmethod
     def generate_csv_content(csv_content_input):
         pass
 
 
 @register_csv_item(GradConst.MD5)
-class CSV_md5(CsvItem):
+class CsvMd5(CsvItem):
+    @staticmethod
     def generate_csv_header(csv_header_input):
         return ["MD5"]
 
+    @staticmethod
     def generate_csv_content(csv_content_input):
         grad = csv_content_input.grad
         tensor_bytes = grad.cpu().detach().float().numpy().tobytes()
@@ -59,7 +79,8 @@ class CSV_md5(CsvItem):
 
 
 @register_csv_item(GradConst.DISTRIBUTION)
-class CSV_distribution(CsvItem):
+class CsvDistribution(CsvItem):
+    @staticmethod
     def generate_csv_header(csv_header_input):
         bounds = csv_header_input.bounds
         intervals = []
@@ -73,6 +94,7 @@ class CSV_distribution(CsvItem):
     
         return intervals
 
+    @staticmethod
     def generate_csv_content(csv_content_input):
         grad = csv_content_input.grad
         bounds = csv_content_input.bounds
@@ -90,40 +112,48 @@ class CSV_distribution(CsvItem):
 
 
 @register_csv_item(GradConst.MAX)
-class CSV_max(CsvItem):
+class CsvMax(CsvItem):
+    @staticmethod
     def generate_csv_header(csv_header_input):
         return ["max"]
 
+    @staticmethod
     def generate_csv_content(csv_content_input):
         grad = csv_content_input.grad
         return [torch.max(grad).cpu().detach().float().numpy().tolist()]
 
 
 @register_csv_item(GradConst.MIN)
-class CSV_max(CsvItem):
+class CsvMin(CsvItem):
+    @staticmethod
     def generate_csv_header(csv_header_input):
         return ["min"]
 
+    @staticmethod
     def generate_csv_content(csv_content_input):
         grad = csv_content_input.grad
         return [torch.min(grad).cpu().detach().float().numpy().tolist()]
 
 
 @register_csv_item(GradConst.NORM)
-class CSV_max(CsvItem):
+class CsvNorm(CsvItem):
+    @staticmethod
     def generate_csv_header(csv_header_input):
         return ["norm"]
 
+    @staticmethod
     def generate_csv_content(csv_content_input):
         grad = csv_content_input.grad
         return [torch.norm(grad).cpu().detach().float().numpy().tolist()]
     
 
 @register_csv_item(GradConst.SHAPE)
-class CSV_shape(CsvItem):
+class CsvShape(CsvItem):
+    @staticmethod
     def generate_csv_header(csv_header_input):
         return ["shape"]
 
+    @staticmethod
     def generate_csv_content(csv_content_input):
         grad = csv_content_input.grad
         return [list(grad.shape)]

msprobe/pytorch/hook_module/__init__.py

@@ -1 +1,16 @@
-from .wrap_functional import remove_dropout
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .wrap_functional import remove_dropout

msprobe/pytorch/hook_module/api_registry.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2022-2023. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
@@ -13,7 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
 
 import torch
 import torch.distributed as dist
@@ -107,7 +105,14 @@ class ApiRegistry:
         if not is_gpu:
             self.set_api_attr(torch_npu, self.torch_npu_ori_attr)
 
-    def initialize_hook(self, hook):
+    def initialize_hook(self, hook, online_run_ut=False):
+        """
+        initialize_hook
+        Args:
+            hook (_type_): initialize_hook
+            online_run_ut (bool): default False, whether online run_ut or not.
+                If online_run_ut is True, the hook will not wrap the aten ops.
+        """
         self.store_ori_attr(torch.Tensor, get_tensor_ops(), self.tensor_ori_attr)
         wrap_tensor.wrap_tensor_ops_and_bind(hook)
         for attr_name in dir(wrap_tensor.HOOKTensor):
@@ -137,7 +142,7 @@ class ApiRegistry:
                     self.npu_distributed_hook_attr[attr_name[5:]] = getattr(wrap_distributed.HOOKDistributedOP,
                                                                             attr_name)
 
-        if torch_version_above_2:
+        if torch_version_above_2 and not online_run_ut:
             self.store_ori_attr(torch.ops.aten, get_aten_ops(), self.aten_ori_attr)
             wrap_aten.wrap_aten_ops_and_bind(hook)
             for attr_name in dir(wrap_aten.HOOKAtenOP):

msprobe/pytorch/hook_module/hook_module.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2019-2020. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
@@ -13,7 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
 
 import functools
 import threading
@@ -58,12 +56,12 @@ class HOOKModule(nn.Module):
                 self.register_forward_hook(forward_hook)
             self.register_backward_hook(backward_hook)
 
-    def __call__(self, *input, **kwargs):
+    def __call__(self, *args, **kwargs):
         changed = False
         if not self.stop_hook:
             HOOKModule.inner_stop_hook[self.current_thread] = True
             changed = True
-        result = self._call_func(*input, **kwargs)
+        result = self._call_func(*args, **kwargs)
         if changed:
             HOOKModule.inner_stop_hook[self.current_thread] = False
         return result
@@ -72,28 +70,28 @@ class HOOKModule(nn.Module):
     def reset_module_stats(cls):
         cls.module_count = {}
 
-    def _call_func(self, *input, **kwargs):
+    def _call_func(self, *args, **kwargs):
         full_backward_hooks, non_full_backward_hooks = [], []
         if len(self._backward_hooks) > 0:
             full_backward_hooks, non_full_backward_hooks = self._get_backward_hooks()
         for hook in self._forward_pre_hooks.values():
-            result_input, result_kwargs = hook(self, input, kwargs)
-            if result_input is not None:
-                if not isinstance(result_input, tuple):
-                    result_input = (result_input,)
-                input = result_input
+            result_args, result_kwargs = hook(self, args, kwargs)
+            if result_args is not None:
+                if not isinstance(result_args, tuple):
+                    result_args = (result_args,)
+                args = result_args
             if result_kwargs is not None:
                 kwargs = result_kwargs
         bw_hook = None
         if len(full_backward_hooks) > 0:
             bw_hook = full_hooks.BackwardHook(self, full_backward_hooks)
-            input = bw_hook.setup_input_hook(input)
+            args = bw_hook.setup_input_hook(args)
         if torch._C._get_tracing_state():
-            result = self._slow_forward(*input, **kwargs)
+            result = self._slow_forward(*args, **kwargs)
         else:
-            result = self.forward(*input, **kwargs)
+            result = self.forward(*args, **kwargs)
         for hook in self._forward_hooks.values():
-            hook_result = hook(self, input, kwargs, result)
+            hook_result = hook(self, args, kwargs, result)
             if hook_result is not None:
                 result = hook_result
         if bw_hook:
@@ -116,5 +114,5 @@ class HOOKModule(nn.Module):
                     wrapper = functools.partial(hook, self)
                     functools.update_wrapper(wrapper, hook)
                     grad_fn.register_hook(wrapper)
-                self._maybe_warn_non_full_backward_hook(input, result, grad_fn)
+                self._maybe_warn_non_full_backward_hook(args, result, grad_fn)
         return result

msprobe/pytorch/hook_module/support_wrap_ops.yaml

@@ -1130,6 +1130,7 @@ torch_npu:
   - npu_prompt_flash_attention
   - npu_lstm
   - npu_apply_adam
+  - npu_apply_adam_w
 
 aten:
   - signbit

msprobe/pytorch/hook_module/utils.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2024-2024. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
@@ -13,7 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
 
 import os
 from msprobe.core.common.file_utils import load_yaml

msprobe/pytorch/hook_module/wrap_aten.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2023-2023. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
@@ -13,7 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
 
 import os
 import torch
@@ -24,6 +22,7 @@ from msprobe.core.common.const import Const
 from msprobe.core.common.file_utils import load_yaml
 from msprobe.pytorch.function_factory import npu_custom_grad_functions
 
+
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, "support_wrap_ops.yaml")
 ops = load_yaml(yaml_path)
@@ -50,6 +49,8 @@ class AtenOPTemplate(HOOKModule):
     def __init__(self, op, hook, need_hook=True):
         if isinstance(op, torch._ops.OpOverloadPacket):
             op_name_ = op._qualified_op_name.split("::")[-1]
+        elif isinstance(op, str):
+            op_name_ = str(op)
         else:
             op_name_ = op.name().split("::")[-1]
             overload_name = op._overloadname
@@ -76,13 +77,13 @@ class AtenOPTemplate(HOOKModule):
 
 
 class AtenOPPacketTemplate():
-    def __init__(self, opPacket, hook):
-        self.opPacket = opPacket
+    def __init__(self, op_packet, hook):
+        self.op_packet = op_packet
         self.hook = hook
 
     def __getattr__(self, key):
         try:
-            attr = getattr(self.opPacket, key)
+            attr = getattr(self.op_packet, key)
         except AttributeError as e:
             raise AttributeError(f"AtenOPPacketTemplate or OpOverloadPacket does not have attribute '{key}'.") from e
         if isinstance(attr, torch._ops.OpOverload):
@@ -92,10 +93,10 @@ class AtenOPPacketTemplate():
 
     @torch_device_guard
     def __call__(self, *args, **kwargs):
-        return AtenOPTemplate(self.opPacket, self.hook)(*args, **kwargs)
+        return AtenOPTemplate(self.op_packet, self.hook)(*args, **kwargs)
 
     def overloads(self):
-        return self.opPacket.overloads()
+        return self.op_packet.overloads()
 
 
 def wrap_aten_op(op, hook):