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

msprobe/mindspore/api_accuracy_checker/data_manager.py

@@ -0,0 +1,264 @@
+# 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
+import csv
+
+from msprobe.core.common.const import Const, CompareConst, MsCompareConst
+from msprobe.core.common.file_utils import FileOpen, create_directory, write_csv, read_csv
+from msprobe.core.common.utils import add_time_as_suffix, MsprobeBaseException
+from msprobe.mindspore.api_accuracy_checker.base_compare_algorithm import compare_algorithms
+from msprobe.core.common.file_utils import check_file_or_directory_path
+from msprobe.mindspore.common.log import logger
+
+
+class ResultCsvEntry:
+    def __init__(self) -> None:
+        self.forward_pass_status = None
+        self.backward_pass_status = None
+        self.forward_err_msg = ""
+        self.backward_err_msg = ""
+        self.overall_err_msg = None
+
+
+def write_csv_header(csv_path, header_func):
+    """如果是第一次写入，则写入 CSV 表头"""
+    header = header_func()  # 获取表头
+    logger.debug(f"Writing CSV header: {header}")
+    write_csv([header], csv_path, mode="a+")
+
+
+def get_result_csv_header():
+    """获取结果 CSV 文件的表头"""
+    return [
+        MsCompareConst.DETAIL_CSV_API_NAME,
+        MsCompareConst.RESULT_CSV_FORWARD_TEST_SUCCESS,
+        MsCompareConst.RESULT_CSV_BACKWARD_TEST_SUCCESS,
+        MsCompareConst.DETAIL_CSV_MESSAGE,
+    ]
+
+
+def get_detail_csv_header():
+    """获取详细 CSV 文件的表头"""
+    detail_csv_header_basic_info = [
+        MsCompareConst.DETAIL_CSV_API_NAME,
+        MsCompareConst.DETAIL_CSV_BENCH_DTYPE,
+        MsCompareConst.DETAIL_CSV_TESTED_DTYPE,
+        MsCompareConst.DETAIL_CSV_SHAPE,
+    ]
+    detail_csv_header_compare_result = list(compare_algorithms.keys())
+    detail_csv_header_status = [
+        MsCompareConst.DETAIL_CSV_PASS_STATUS,
+        MsCompareConst.DETAIL_CSV_MESSAGE,
+    ]
+    return detail_csv_header_basic_info + detail_csv_header_compare_result + detail_csv_header_status
+
+
+def check_csv_header(headers, required_constants, csv_path):
+    """校验 CSV 文件表头是否包含所有必需的常量"""
+    missing_constants = [const for const in required_constants if not any(const in header for header in headers)]
+
+    if missing_constants:
+        raise MsprobeBaseException(
+            MsprobeBaseException.MISSING_HEADER_ERROR,
+            f"{csv_path} 缺少以下必需的表头字段: {missing_constants}"
+        )
+
+
+class DataManager:
+    def __init__(self, csv_dir, result_csv_path):
+        self.results = {}
+        self.is_first_write = True  # 标记用于添加表头
+        self.csv_dir = csv_dir
+        self.api_names_set = set()  # 存储已经出现的 API 名称的集合
+        # 如果传入了 result_csv_path，则启用断点续检
+        if result_csv_path:
+            self.resume_from_last_csv(result_csv_path)
+            self.initialize_api_names_set(result_csv_path)
+        else:
+            # 默认情况下，设置输出路径为空，等待首次写入时初始化
+            self.result_out_path = os.path.join(self.csv_dir, add_time_as_suffix(MsCompareConst.RESULT_CSV_FILE_NAME))
+            self.detail_out_path = os.path.join(
+                self.csv_dir,
+                os.path.basename(self.result_out_path).replace("result", "details")
+            )
+
+            if self.detail_out_path and os.path.exists(self.detail_out_path):
+                check_file_or_directory_path(self.detail_out_path)
+
+            if self.result_out_path and os.path.exists(self.result_out_path):
+                check_file_or_directory_path(self.result_out_path)
+
+    def initialize_api_names_set(self, result_csv_path):
+        """读取现有的 CSV 文件并存储已经出现的 API 名称到集合中"""
+        # 使用新的 read_csv 函数读取数据
+        csv_data = read_csv(result_csv_path, as_pd=False)
+
+        # 读取标题行
+        headers = csv_data[0] if csv_data else []  # 如果文件为空，则 headers 会为空
+
+        # 使用提取的表头校验函数
+        if check_csv_header(headers, get_result_csv_header(), result_csv_path):
+
+            # 获取 "API Name" 列的索引
+            api_name_index = None
+            for i, header in enumerate(headers):
+                if MsCompareConst.DETAIL_CSV_API_NAME in header:  # CSV 文件的标题行包含了字节顺序标记,所以使用通过包含方式来查找
+                    api_name_index = i
+                    break
+
+            if api_name_index is None:
+                logger.warning(f"{result_csv_path} No column contains 'API Name'.")
+                return
+
+            # 读取每一行的 API 名称
+            for row in csv_data[1:]:  # 跳过标题行，从第二行开始
+                if row and len(row) > api_name_index:
+                    api_name = row[api_name_index]
+                    if api_name:
+                        self.api_names_set.add(api_name)
+
+            logger.debug(f"Initialized API names set from existing CSV: {self.api_names_set}")
+
+    def is_unique_api(self, api_name):
+        """检查 API 名称是否唯一，如果已经存在则返回 False，否则加入集合并返回 True"""
+        if api_name in self.api_names_set:
+            return False
+        self.api_names_set.add(api_name)
+        return True
+
+    def resume_from_last_csv(self, result_csv_path):
+        """从上次运行的 result_csv_path 恢复断点"""
+        # 获取上次的目录路径
+        last_dir = os.path.dirname(result_csv_path)
+
+        # 设置当前目录和输出路径，确保在首次写入时使用
+        self.csv_dir = last_dir
+        self.detail_out_path = os.path.join(last_dir, os.path.basename(result_csv_path).replace("result", "details"))
+        if self.detail_out_path and os.path.exists(self.detail_out_path):
+            check_file_or_directory_path(self.detail_out_path)
+        self.result_out_path = result_csv_path
+        self.is_first_write = False
+
+    def save_results(self, api_name_str):
+        if self.is_first_write:
+            # 直接写入表头
+            logger.info("Writing CSV headers for the first time.")
+            write_csv_header(self.detail_out_path, get_detail_csv_header)
+            write_csv_header(self.result_out_path, get_result_csv_header)
+            self.is_first_write = False  # 写入后标记为 False，避免重复写入表头
+
+        """写入详细输出和结果摘要并清理结果"""
+        logger.debug("Starting to write detailed output to CSV.")
+        self.to_detail_csv(self.detail_out_path)
+        logger.debug(f"Detailed output for {api_name_str} written to {self.detail_out_path}.")
+
+        logger.debug("Starting to write result summary to CSV.")
+        self.to_result_csv(self.result_out_path)
+        logger.debug(f"Result summary for {api_name_str} written to {self.result_out_path}.")
+
+        # 清理记录，准备下一次调用
+        self.clear_results()
+
+    def record(self, output_list):
+        if output_list is None:
+            return
+        for output in output_list:
+            api_real_name, forward_or_backward, basic_info, compare_result_dict = output
+            key = (api_real_name, forward_or_backward)
+            if key not in self.results:
+                self.results[key] = []
+            self.results[key].append((basic_info, compare_result_dict))
+            logger.debug(f"Updated self.results for key {key}: {self.results[key]}")
+        logger.debug(f"Complete self.results after recording: {self.results}")
+
+    def clear_results(self):
+        """清空 self.results 数据"""
+        logger.debug("Clearing self.results data.")
+        self.results.clear()
+
+    def to_detail_csv(self, csv_path):
+        logger.debug("Preparing detail CSV headers and rows.")
+        detail_csv = []
+
+        detail_csv_header_compare_result = list(compare_algorithms.keys())
+
+        for _, results in self.results.items():
+            for res in results:
+                basic_info, compare_result_dict = res
+                csv_row_basic_info = [
+                    basic_info.api_name,
+                    basic_info.bench_dtype,
+                    basic_info.tested_dtype,
+                    basic_info.shape
+                ]
+                csv_row_compare_result = [
+                    compare_result_dict.get(algorithm_name).compare_value
+                    for algorithm_name in detail_csv_header_compare_result
+                ]
+                csv_row_status = [basic_info.status, basic_info.err_msg]
+                csv_row = csv_row_basic_info + csv_row_compare_result + csv_row_status
+                detail_csv.append(csv_row)
+                logger.debug(f"Detail CSV row added: {csv_row}")
+
+        logger.debug(f"Writing detail CSV to {csv_path}.")
+        write_csv(detail_csv, csv_path, mode="a+")
+        logger.debug(f"Detail CSV written successfully to {csv_path}.")
+
+    def to_result_csv(self, csv_path):
+        logger.debug("Preparing result CSV data.")
+        result_csv = []
+
+        result_csv_dict = {}
+        for key, results in self.results.items():
+            api_real_name, forward_or_backward = key
+            pass_status = CompareConst.PASS
+            overall_err_msg = ""
+
+            for res in results:
+                basic_info, _ = res
+                if basic_info.status != CompareConst.PASS:
+                    pass_status = CompareConst.ERROR
+                overall_err_msg += basic_info.err_msg
+
+            overall_err_msg = "" if pass_status == CompareConst.PASS else overall_err_msg
+
+            if api_real_name not in result_csv_dict:
+                result_csv_dict[api_real_name] = ResultCsvEntry()
+            if forward_or_backward == Const.FORWARD:
+                result_csv_dict[api_real_name].forward_pass_status = pass_status
+                result_csv_dict[api_real_name].forward_err_msg = overall_err_msg
+            else:
+                result_csv_dict[api_real_name].backward_pass_status = pass_status
+                result_csv_dict[api_real_name].backward_err_msg = overall_err_msg
+
+        for api_name, entry in result_csv_dict.items():
+            overall_err_msg = "" if (entry.forward_pass_status == CompareConst.PASS and
+                                     entry.backward_pass_status == CompareConst.PASS) else \
+                entry.forward_err_msg + entry.backward_err_msg
+            row = [
+                api_name,
+                entry.forward_pass_status,
+                entry.backward_pass_status,
+                overall_err_msg
+            ]
+            result_csv.append(row)
+            logger.debug(f"Result CSV row added: {row}")
+
+        write_csv(result_csv, csv_path, mode="a+")
+        logger.debug(f"Result CSV written successfully to {csv_path}.")
+
+        # 设置标记为 False，防止后续重复添加表头
+        self.is_first_write = False

msprobe/mindspore/api_accuracy_checker/main.py

@@ -1,9 +1,33 @@
+# 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.mindspore.api_accuracy_checker.api_accuracy_checker import ApiAccuracyChecker
 
+from msprobe.mindspore.api_accuracy_checker.multi_api_accuracy_checker import MultiApiAccuracyChecker
+
+from msprobe.mindspore.api_accuracy_checker.cmd_parser import check_args
+
 
 def api_checker_main(args):
-    api_accuracy_checker = ApiAccuracyChecker()
+    check_args(args)
+    api_accuracy_checker = ApiAccuracyChecker(args)
+    api_accuracy_checker.parse(args.api_info_file)
+    api_accuracy_checker.run_and_compare()
+
+def mul_api_checker_main(args):
+    check_args(args)
+    api_accuracy_checker = MultiApiAccuracyChecker(args)
     api_accuracy_checker.parse(args.api_info_file)
     api_accuracy_checker.run_and_compare()
-    api_accuracy_checker.to_detail_csv(args.out_path)
-    api_accuracy_checker.to_result_csv(args.out_path)

msprobe/mindspore/api_accuracy_checker/multi_api_accuracy_checker.py

@@ -0,0 +1,206 @@
+# 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 multiprocessing
+from multiprocessing import Manager
+import os
+import signal
+import sys
+import time
+
+# 第三方库导入
+from mindspore import context
+import numpy as np
+from tqdm import tqdm
+
+# 本地应用/库特定导入
+from msprobe.core.common.const import Const, CompareConst, MsCompareConst
+from msprobe.mindspore.api_accuracy_checker.api_accuracy_checker import ApiAccuracyChecker, BasicInfoAndStatus
+from msprobe.mindspore.api_accuracy_checker.multi_data_manager import MultiDataManager
+from msprobe.mindspore.common.log import logger
+
+
+class MultiApiAccuracyChecker(ApiAccuracyChecker):
+    def __init__(self, args):
+        # 可以添加 MultiApiAccuracyChecker 特有的属性或方法
+        self.api_infos = dict()
+
+        # 使用 Manager 创建共享变量，确保进程间的同步
+        self.manager = Manager()
+        self.is_first_write = self.manager.Value('b', True)  # 创建共享变量
+
+        # 初始化 DataManager 时传入共享的 is_first_write
+        self.multi_data_manager = MultiDataManager(args.out_path, args.result_csv_path, self.is_first_write)
+
+        self.args = args  # 将 args 保存为类的属性
+
+        # 初始化一个属性来存储当前的设备ID（用于日志中显示）
+        self.current_device_id = None
+
+    def process_on_device(self, device_id, api_infos, progress_queue):
+        """
+        在特定设备上处理一部分API。
+
+        参数:
+            device_id (int): 要使用的设备ID。
+            api_infos (list): 包含API名称和对应信息的元组列表。
+            progress_queue (multiprocessing.Queue): 用于通信进度更新的队列。
+        """
+
+        # 设置当前设备ID
+        self.current_device_id = device_id
+
+        # 设置 MindSpore context 的 device_id
+        context.set_context(device_id=device_id)
+
+        # 遍历当前进程分配的任务
+        for _, (api_name_str, api_info) in enumerate(api_infos):
+            logger.debug(f"Processing API: {api_name_str}, Device: {device_id}")
+
+            if not self.multi_data_manager.is_unique_api(api_name_str):
+                logger.debug(f"API {api_name_str} is not unique, skipping.")
+                progress_queue.put(1)
+                continue
+
+            # 处理前向
+            forward_output_list = self.process_forward(api_name_str, api_info)
+            if forward_output_list is not Const.EXCEPTION_NONE:
+                self.multi_data_manager.record(forward_output_list)
+
+            # 处理反向
+            backward_output_list = self.process_backward(api_name_str, api_info)
+            if backward_output_list is not Const.EXCEPTION_NONE:
+                self.multi_data_manager.record(backward_output_list)
+
+            # 保存结果
+            self.multi_data_manager.save_results(api_name_str)
+            progress_queue.put(1)  # 更新进度
+
+    def run_and_compare(self):
+        # 获取要使用的设备ID列表
+        device_ids = self.args.device_id
+
+        # 按设备数划分要处理的 API 项
+        partitioned_api_infos = list(self.api_infos.items())
+
+        # 在主进程中进行交叉任务切分（基于取模的方式）
+        partitioned_api_infos_split = [[] for _ in range(len(device_ids))]
+        for idx, api_info in enumerate(partitioned_api_infos):
+            device_index = idx % len(device_ids)  # 使用取模方法分配任务
+            partitioned_api_infos_split[device_index].append(api_info)
+
+        # 创建一个共享进度队列
+        progress_queue = multiprocessing.Queue()
+
+        # 进度条
+        total_tasks = len(partitioned_api_infos)  # 计算总任务数
+        with tqdm(total=total_tasks, desc="Total Progress", ncols=100) as pbar:
+            # 创建多进程
+            processes = []
+            for index, device_id in enumerate(device_ids):
+                process = multiprocessing.Process(target=self.process_on_device,
+                                                  args=(device_id, partitioned_api_infos_split[index], progress_queue))
+                processes.append(process)
+                process.start()
+
+            # 主进程更新进度条
+            completed_tasks = 0
+            while completed_tasks < total_tasks:
+                try:
+                    completed_tasks += progress_queue.get(timeout=Const.PROGRESS_TIMEOUT)  # 设置超时时间（秒）
+                    pbar.update(1)
+                except multiprocessing.queues.Empty:
+                    logger.error("Timeout while waiting for progress updates. Skipping remaining tasks.")
+                    break
+
+                # 检查子进程状态
+                for process in processes:
+                    if not process.is_alive():
+                        if process.exitcode != 0:
+                            logger.error(f"Process {process.pid} exited with code {process.exitcode}.")
+                            total_tasks -= len(partitioned_api_infos_split[processes.index(process)])
+                        processes.remove(process)
+
+            # 确保所有子进程完成或终止
+            for process in processes:
+                process.join(timeout=Const.PROGRESS_TIMEOUT)
+                if process.is_alive():
+                    logger.error(f"Process {process.pid} did not terminate. Forcing termination.")
+                    process.terminate()
+
+    def process_forward(self, api_name_str, api_info):
+        """
+        Overrides the parent class's process_forward method to log the device ID when exceptions occur.
+
+        Parameters:
+            api_name_str (str): The name of the API.
+            api_info (object): The API information object.
+
+        Returns:
+            list or None: The forward output list or None if an error occurs.
+        """
+        if not api_info.check_forward_info():
+            logger.debug(
+                f"[Device {self.current_device_id}] API: {api_name_str} lacks forward information, skipping forward check.")
+            return Const.EXCEPTION_NONE
+
+        try:
+            forward_inputs_aggregation = self.prepare_api_input_aggregation(api_info, Const.FORWARD)
+        except Exception as e:
+            logger.warning(
+                f"[Device {self.current_device_id}] Exception occurred while getting forward API inputs for {api_name_str}. Skipping forward check. Detailed exception information: {e}.")
+            return Const.EXCEPTION_NONE
+
+        forward_output_list = None
+        try:
+            forward_output_list = self.run_and_compare_helper(api_info, api_name_str, forward_inputs_aggregation,
+                                                              Const.FORWARD)
+        except Exception as e:
+            logger.warning(
+                f"[Device {self.current_device_id}] Exception occurred while running and comparing {api_name_str} forward API. Detailed exception information: {e}.")
+        return forward_output_list
+
+    def process_backward(self, api_name_str, api_info):
+        """
+        Overrides the parent class's process_backward method to log the device ID when exceptions occur.
+
+        Parameters:
+            api_name_str (str): The name of the API.
+            api_info (object): The API information object.
+
+        Returns:
+            list or None: The backward output list or None if an error occurs.
+        """
+        if not api_info.check_backward_info():
+            logger.debug(
+                f"[Device {self.current_device_id}] API: {api_name_str} lacks backward information, skipping backward check.")
+            return Const.EXCEPTION_NONE
+
+        try:
+            backward_inputs_aggregation = self.prepare_api_input_aggregation(api_info, Const.BACKWARD)
+        except Exception as e:
+            logger.warning(
+                f"[Device {self.current_device_id}] Exception occurred while getting backward API inputs for {api_name_str}. Skipping backward check. Detailed exception information: {e}.")
+            return Const.EXCEPTION_NONE
+
+        backward_output_list = None
+        try:
+            backward_output_list = self.run_and_compare_helper(api_info, api_name_str, backward_inputs_aggregation,
+                                                               Const.BACKWARD)
+        except Exception as e:
+            logger.warning(
+                f"[Device {self.current_device_id}] Exception occurred while running and comparing {api_name_str} backward API. Detailed exception information: {e}.")
+        return backward_output_list

msprobe/mindspore/api_accuracy_checker/multi_data_manager.py

@@ -0,0 +1,58 @@
+# 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 multiprocessing
+import os
+
+from msprobe.mindspore.api_accuracy_checker.data_manager import DataManager, ResultCsvEntry, write_csv_header, get_result_csv_header, get_detail_csv_header, check_csv_header
+from msprobe.mindspore.common.log import logger
+
+
+class MultiDataManager(DataManager):
+    def __init__(self, csv_dir, result_csv_path, shared_is_first_write):
+        super().__init__(csv_dir, result_csv_path)
+
+        # 使用共享的 is_first_write 变量来控制表头写入
+        self.shared_is_first_write = shared_is_first_write
+        # 创建锁对象，确保线程安全
+        self.lock = multiprocessing.Lock()
+
+    def save_results(self, api_name_str):
+        """保存结果，线程安全操作"""
+
+        with self.lock:  # 确保保存操作不会被多个进程同时进行
+            if self.is_first_write and self.shared_is_first_write.value:
+                self.shared_is_first_write.value = False
+                self.is_first_write = False  # 写入后标记为 False，避免重复写入表头
+                # 直接写入表头
+                logger.info("Writing CSV headers for the first time.")
+                write_csv_header(self.detail_out_path, get_detail_csv_header)
+                write_csv_header(self.result_out_path, get_result_csv_header)
+
+            """写入详细输出和结果摘要并清理结果"""
+            self.to_detail_csv(self.detail_out_path)
+            logger.debug(f"Detailed output for {api_name_str} written to {self.detail_out_path}.")
+
+            self.to_result_csv(self.result_out_path)
+            logger.debug(f"Result summary for {api_name_str} written to {self.result_out_path}.")
+
+            # 清理记录，准备下一次调用
+            self.clear_results()
+
+    def clear_results(self):
+        """清空 self.results 数据，线程安全操作"""
+        logger.debug("Clearing results data.")
+        self.results.clear()

msprobe/mindspore/api_accuracy_checker/type_mapping.py

@@ -1,7 +1,23 @@
-from mindspore.common import dtype as mstype
-import numpy as np
+# 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 mindspore
+import numpy as np
 import torch
+from mindspore._c_expression import typing
+from mindspore.common import dtype as mstype
 
 INT8 = "Int8"
 UINT8 = "UInt8"
@@ -18,7 +34,6 @@ BOOL = "Bool"
 BFLOAT16 = "BFloat16"
 INT4 = "Int4"
 
-
 dtype_str_to_ms_dtype = {
     INT8: mstype.int8,
     UINT8: mstype.uint8,
@@ -37,7 +52,6 @@ dtype_str_to_ms_dtype = {
 }
 ms_dtype_to_dtype_str = {value: key for key, value in dtype_str_to_ms_dtype.items()}
 
-
 dtype_str_to_np_dtype = {
     INT8: np.int8,
     UINT8: np.uint8,
@@ -75,6 +89,8 @@ FLOAT_TYPE_STR = "float"
 SLICE_TYPE_STR = "slice"
 TUPLE_TYPE_STR = "tuple"
 STR_TYPE_STR = "str"
+MINDSPORE_DTYPE_TYPE_STR = "mindspore.dtype"
+TORCH_DTYPE_TYPE_STR = "torch.dtype"
 
 api_info_type_str_to_type = {
     MINDSPORE_TENSOR_TYPE_STR: mindspore.Tensor,
@@ -83,6 +99,7 @@ api_info_type_str_to_type = {
     FLOAT_TYPE_STR: float,
     SLICE_TYPE_STR: slice,
     STR_TYPE_STR: str,
+    MINDSPORE_DTYPE_TYPE_STR: typing.Type,
 }
 type_to_api_info_type_str = {value: key for key, value in api_info_type_str_to_type.items()}
 
@@ -111,4 +128,4 @@ uint_dtype_str_list = [
     UINT16,
     UINT32,
     UINT64,
-]
+]