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

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/server.py

@@ -1,218 +1,219 @@
-import os.path
-import struct
-import hashlib
-import time
-import io
-from threading import Thread
-
-from OpenSSL import SSL
-from twisted.internet import ssl, reactor, protocol, endpoints
-
-from msprobe.pytorch.common.utils import logger
-from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.ssl_config import cipher_list
-
-
-class TCPServer:
-    def __init__(self, port, shared_queue, check_sum=False, tls_path=None) -> None:
-        self.port = port
-        self.shared_queue = shared_queue
-        self.check_sum = check_sum
-        self.tls_path = tls_path
-        self.factory = MessageServerFactory()
-        self.reactor_thread = None
-
-    @staticmethod
-    def run_reactor():
-        reactor.run(installSignalHandlers=False)
-
-    def start(self):
-        self.factory.protocol = self.build_protocol
-
-        if self.tls_path:
-            server_key = os.path.join(self.tls_path, "server.key")
-            server_crt = os.path.join(self.tls_path, "server.crt")
-            server_context_factory = ssl.DefaultOpenSSLContextFactory(server_key, server_crt, SSL.TLSv1_2_METHOD)
-            server_context_ = server_context_factory.getContext()
-            server_context_.set_cipher_list(cipher_list)
-            server_context_.set_options(SSL.OP_NO_RENEGOTIATION)
-            endpoint = endpoints.SSL4ServerEndpoint(reactor, self.port, server_context_factory)
-        else:
-            endpoint = endpoints.TCP4ServerEndpoint(reactor, self.port)
-        endpoint.listen(self.factory)
-        self.reactor_thread = Thread(target=self.run_reactor, daemon=True)
-        self.reactor_thread.start()
-
-    def is_running(self):
-        return not self.factory.is_all_connection_closed()
-
-    def stop(self):
-        self.factory.doStop()
-        reactor.callFromThread(reactor.sigInt, 2)
-        self.reactor_thread.join()
-
-    def build_protocol(self):
-        return ServerProtocol(self.shared_queue, self.check_sum)
-
-
-class ServerProtocol(protocol.Protocol):
-    ACK_SUCCESS = b"OK___"
-    ACK_ERROR = b"ERROR"
-    ACK_BUSY = b"BUSY_"
-    ACK_STOP = b"STOP_"
-    ACK_STOP_CONFIRM = b"OVER_"
-    ACK_KILL_PROCESS = b"KILL_"
-
-    def __init__(self, shared_queue, check_sum=False):
-        self.start_time = None
-        self.buffer = io.BytesIO()
-        self.consumer_queue = shared_queue
-        self.check_sum = check_sum
-        self.length_width = 8
-        self.md5_width = 32
-        self.obj_length = None
-        self.tell = 0
-        self.obj_md5 = None
-        self.obj_body = None
-        self.sequence_number = -1
-        self.rank = -1
-        self.step = -1
-        self.sequence_number_dict = dict()
-
-    def connectionMade(self):
-        self.buffer = io.BytesIO()
-        self.obj_length = None
-        self.tell = 0
-        self.obj_md5 = None
-        self.obj_body = None
-        self.factory.transport_dict[self.transport] = 1
-        self.factory.transport_list.append(self.transport)
-        logger.info(f"Connected to {self.transport.getPeer()} successfully.")
-
-    def connectionLost(self, reason):
-        self.factory.transport_dict.pop(self.transport, None)
-        if len(self.factory.transport_dict) == 0:
-            self.consumer_queue.put(self.ACK_KILL_PROCESS)
-
-        logger.info(f"Lost connection with {self.transport.getPeer()}. Reason is: {reason} 与客户端 断开连接, "
-                    f"current connection number is: {len(self.factory.transport_dict)}")
-
-    def send_ack(self, ack_info):
-        ack_message = b"".join([
-            ack_info,
-            self.sequence_number.to_bytes(8, byteorder='big'),
-            self.rank.to_bytes(8, byteorder='big'),
-            self.step.to_bytes(8, byteorder='big')
-        ])
-        self.transport.write(ack_message)
-
-    def post_process(self):
-        send_busy_ack = False
-        while self.consumer_queue.full():
-            if not send_busy_ack:
-                self.send_ack(self.ACK_BUSY)
-                logger.debug("sending BUSY ACK")
-            send_busy_ack = True
-            time.sleep(0.1)
-
-        obj_key = str(self.sequence_number) + "_" + str(self.rank) + "_" + str(self.step)
-
-        recv_md5 = hashlib.md5(self.obj_body).hexdigest()
-        if self.check_sum and recv_md5 != self.obj_md5:
-            # when needs check md5 and check no pass, indicates received data error, send b"ERROR" to client.
-            logger.debug(f"Error:接收数据有问题，流水号{self.sequence_number}, expected {self.obj_md5}, but get {recv_md5}")
-            self.send_ack(self.ACK_ERROR)
-        else:
-            if self.obj_body == self.ACK_STOP:
-                self.handle_with_stop()
-            else:
-                self.send_ack(self.ACK_SUCCESS)
-            if obj_key in self.sequence_number_dict:
-                logger.debug(f"这是一次异常的重传，可以忽略。 {obj_key}, {self.sequence_number_dict}")
-            else:
-                self.sequence_number_dict[obj_key] = self.obj_md5
-                self.consumer_queue.put(self.obj_body, block=True)
-
-        self.reset_env()
-        finish_time = time.time()
-        logger.debug(f"finish_time: {finish_time - self.start_time}")
-
-    def handle_with_stop(self):
-        logger.debug(f"接收到停止传输信号 TCP{self.transport.getPeer()}")
-        self.send_ack(self.ACK_STOP_CONFIRM)
-        if len(self.factory.transport_dict) == 0:
-            _rank, _step, _sequence_number = 0, 0, 100000000
-            ack_kill = self.ACK_KILL_PROCESS + \
-                       _sequence_number.to_bytes(8, byteorder='big') + \
-                       _rank.to_bytes(8, byteorder='big') + \
-                       _step.to_bytes(8, byteorder='big')
-            for trans in self.factory.transport_list:
-                trans.write(ack_kill)
-            logger.debug(f"发送KILL信息给{self.transport.getPeer()}")
-            self.consumer_queue.put(self.ACK_KILL_PROCESS)
-            time.sleep(2)
-
-    def reset_env(self):
-        self.obj_length = None
-        self.sequence_number = -1
-        self.rank = -1
-        self.step = -1
-        self.obj_md5 = None
-        self.obj_body = None
-
-    def dataReceived(self, data):
-        self.buffer.seek(0, 2)
-        self.buffer.write(data)
-        self.buffer.seek(self.tell)
-
-        # The first data packet is packet header, it contains obj_length, sequence_number, rank, step
-        if self.obj_length is None and len(self.buffer.getvalue()) >= self.length_width * 4:
-            self.start_time = time.time()
-            self.obj_length = struct.unpack('!Q', self.buffer.read(self.length_width))[0]
-            self.sequence_number = struct.unpack('!Q', self.buffer.read(self.length_width))[0]
-            self.rank = struct.unpack('!Q', self.buffer.read(self.length_width))[0]
-            self.step = struct.unpack('!Q', self.buffer.read(self.length_width))[0]
-            self.tell += self.length_width * 4
-            logger.debug(
-                f"流水号: {self.sequence_number}; RANK: {self.rank}; STEP: {self.step}; Length: {self.obj_length}")
-
-        # If needs check md5 but not parse md5 yet, read 32b md5 values
-        check_sum_and_md5 = (self.check_sum
-                             and self.obj_length is not None
-                             and self.obj_md5 is None
-                             and len(self.buffer.getvalue()) - self.tell >= self.md5_width)
-        if check_sum_and_md5:
-            self.obj_md5 = self.buffer.read(self.md5_width).decode()
-            self.tell += self.md5_width
-            logger.debug(f"MD5: {self.obj_md5}")
-
-        current_length = len(self.buffer.getvalue()) - self.tell
-        if self.obj_length is not None and 0 < self.obj_length <= current_length:
-            # Current api data receive finished
-            self.obj_body = self.buffer.read(self.obj_length)
-
-            self.tell += self.obj_length
-            self.buffer = io.BytesIO(self.buffer.getvalue()[self.tell:])
-            self.buffer.seek(0)
-            self.tell = 0
-            recv_data_time = time.time()
-            logger.debug(f"self.sequence_number {self.sequence_number} "
-                         f"recv_data_time {recv_data_time - self.start_time}")
-
-            if self.obj_body == self.ACK_STOP:
-                # Indicates the current TCP link receives a STOP signal and remove from the transport_dict
-                _transport = self.factory.transport_dict.pop(self.transport, None)
-                logger.debug(f"接收到b'STOP_' self.sequence_number {self.sequence_number} ")
-            self.post_process()
-
-
-class MessageServerFactory(protocol.ServerFactory):
-    def __init__(self) -> None:
-        """
-        transport_dict: links that have not completed data transmission.
-        transport_list: Records all TCP links. Appends TCP link to the transport list when a new TCP link is established.
-        """
-        self.transport_dict = {}
-        self.transport_list = []
-
-    def is_all_connection_closed(self):
-        return len(self.transport_dict) == 0
+import os.path
+import struct
+import hashlib
+import time
+import io
+from threading import Thread
+
+from twisted.internet import reactor, protocol, endpoints
+
+from msprobe.pytorch.common.utils import logger
+from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.ssl_config import cipher_list
+
+
+class TCPServer:
+    def __init__(self, port, shared_queue, check_sum=False, tls_path=None) -> None:
+        self.port = port
+        self.shared_queue = shared_queue
+        self.check_sum = check_sum
+        self.tls_path = tls_path
+        self.factory = MessageServerFactory()
+        self.reactor_thread = None
+
+    @staticmethod
+    def run_reactor():
+        reactor.run(installSignalHandlers=False)
+
+    def start(self):
+        self.factory.protocol = self.build_protocol
+
+        if self.tls_path:
+            from OpenSSL import SSL
+            from twisted.internet import ssl
+            server_key = os.path.join(self.tls_path, "server.key")
+            server_crt = os.path.join(self.tls_path, "server.crt")
+            server_context_factory = ssl.DefaultOpenSSLContextFactory(server_key, server_crt, SSL.TLSv1_2_METHOD)
+            server_context_ = server_context_factory.getContext()
+            server_context_.set_cipher_list(cipher_list)
+            server_context_.set_options(SSL.OP_NO_RENEGOTIATION)
+            endpoint = endpoints.SSL4ServerEndpoint(reactor, self.port, server_context_factory)
+        else:
+            endpoint = endpoints.TCP4ServerEndpoint(reactor, self.port)
+        endpoint.listen(self.factory)
+        self.reactor_thread = Thread(target=self.run_reactor, daemon=True)
+        self.reactor_thread.start()
+
+    def is_running(self):
+        return not self.factory.is_all_connection_closed()
+
+    def stop(self):
+        self.factory.doStop()
+        reactor.callFromThread(reactor.sigInt, 2)
+        self.reactor_thread.join()
+
+    def build_protocol(self):
+        return ServerProtocol(self.shared_queue, self.check_sum)
+
+
+class ServerProtocol(protocol.Protocol):
+    ACK_SUCCESS = b"OK___"
+    ACK_ERROR = b"ERROR"
+    ACK_BUSY = b"BUSY_"
+    ACK_STOP = b"STOP_"
+    ACK_STOP_CONFIRM = b"OVER_"
+    ACK_KILL_PROCESS = b"KILL_"
+
+    def __init__(self, shared_queue, check_sum=False):
+        self.start_time = None
+        self.buffer = io.BytesIO()
+        self.consumer_queue = shared_queue
+        self.check_sum = check_sum
+        self.length_width = 8
+        self.md5_width = 32
+        self.obj_length = None
+        self.tell = 0
+        self.obj_md5 = None
+        self.obj_body = None
+        self.sequence_number = -1
+        self.rank = -1
+        self.step = -1
+        self.sequence_number_dict = dict()
+
+    def connectionMade(self):
+        self.buffer = io.BytesIO()
+        self.obj_length = None
+        self.tell = 0
+        self.obj_md5 = None
+        self.obj_body = None
+        self.factory.transport_dict[self.transport] = 1
+        self.factory.transport_list.append(self.transport)
+        logger.info(f"Connected to {self.transport.getPeer()} successfully.")
+
+    def connectionLost(self, reason):
+        self.factory.transport_dict.pop(self.transport, None)
+        if len(self.factory.transport_dict) == 0:
+            self.consumer_queue.put(self.ACK_KILL_PROCESS)
+
+        logger.info(f"Lost connection with {self.transport.getPeer()}. Reason is: {reason} 与客户端 断开连接, "
+                    f"current connection number is: {len(self.factory.transport_dict)}")
+
+    def send_ack(self, ack_info):
+        ack_message = b"".join([
+            ack_info,
+            self.sequence_number.to_bytes(8, byteorder='big'),
+            self.rank.to_bytes(8, byteorder='big'),
+            self.step.to_bytes(8, byteorder='big')
+        ])
+        self.transport.write(ack_message)
+
+    def post_process(self):
+        send_busy_ack = False
+        while self.consumer_queue.full():
+            if not send_busy_ack:
+                self.send_ack(self.ACK_BUSY)
+                logger.debug("sending BUSY ACK")
+            send_busy_ack = True
+            time.sleep(0.1)
+
+        obj_key = str(self.sequence_number) + "_" + str(self.rank) + "_" + str(self.step)
+
+        recv_md5 = hashlib.md5(self.obj_body).hexdigest()
+        if self.check_sum and recv_md5 != self.obj_md5:
+            # when needs check md5 and check no pass, indicates received data error, send b"ERROR" to client.
+            logger.debug(f"Error:接收数据有问题，流水号{self.sequence_number}, expected {self.obj_md5}, but get {recv_md5}")
+            self.send_ack(self.ACK_ERROR)
+        else:
+            if self.obj_body == self.ACK_STOP:
+                self.handle_with_stop()
+            else:
+                self.send_ack(self.ACK_SUCCESS)
+            if obj_key in self.sequence_number_dict:
+                logger.debug(f"这是一次异常的重传，可以忽略。 {obj_key}, {self.sequence_number_dict}")
+            else:
+                self.sequence_number_dict[obj_key] = self.obj_md5
+                self.consumer_queue.put(self.obj_body, block=True)
+
+        self.reset_env()
+        finish_time = time.time()
+        logger.debug(f"finish_time: {finish_time - self.start_time}")
+
+    def handle_with_stop(self):
+        logger.debug(f"接收到停止传输信号 TCP{self.transport.getPeer()}")
+        self.send_ack(self.ACK_STOP_CONFIRM)
+        if len(self.factory.transport_dict) == 0:
+            _rank, _step, _sequence_number = 0, 0, 100000000
+            ack_kill = self.ACK_KILL_PROCESS + \
+                       _sequence_number.to_bytes(8, byteorder='big') + \
+                       _rank.to_bytes(8, byteorder='big') + \
+                       _step.to_bytes(8, byteorder='big')
+            for trans in self.factory.transport_list:
+                trans.write(ack_kill)
+            logger.debug(f"发送KILL信息给{self.transport.getPeer()}")
+            self.consumer_queue.put(self.ACK_KILL_PROCESS)
+            time.sleep(2)
+
+    def reset_env(self):
+        self.obj_length = None
+        self.sequence_number = -1
+        self.rank = -1
+        self.step = -1
+        self.obj_md5 = None
+        self.obj_body = None
+
+    def dataReceived(self, data):
+        self.buffer.seek(0, 2)
+        self.buffer.write(data)
+        self.buffer.seek(self.tell)
+
+        # The first data packet is packet header, it contains obj_length, sequence_number, rank, step
+        if self.obj_length is None and len(self.buffer.getvalue()) >= self.length_width * 4:
+            self.start_time = time.time()
+            self.obj_length = struct.unpack('!Q', self.buffer.read(self.length_width))[0]
+            self.sequence_number = struct.unpack('!Q', self.buffer.read(self.length_width))[0]
+            self.rank = struct.unpack('!Q', self.buffer.read(self.length_width))[0]
+            self.step = struct.unpack('!Q', self.buffer.read(self.length_width))[0]
+            self.tell += self.length_width * 4
+            logger.debug(
+                f"流水号: {self.sequence_number}; RANK: {self.rank}; STEP: {self.step}; Length: {self.obj_length}")
+
+        # If needs check md5 but not parse md5 yet, read 32b md5 values
+        check_sum_and_md5 = (self.check_sum
+                             and self.obj_length is not None
+                             and self.obj_md5 is None
+                             and len(self.buffer.getvalue()) - self.tell >= self.md5_width)
+        if check_sum_and_md5:
+            self.obj_md5 = self.buffer.read(self.md5_width).decode()
+            self.tell += self.md5_width
+            logger.debug(f"MD5: {self.obj_md5}")
+
+        current_length = len(self.buffer.getvalue()) - self.tell
+        if self.obj_length is not None and 0 < self.obj_length <= current_length:
+            # Current api data receive finished
+            self.obj_body = self.buffer.read(self.obj_length)
+
+            self.tell += self.obj_length
+            self.buffer = io.BytesIO(self.buffer.getvalue()[self.tell:])
+            self.buffer.seek(0)
+            self.tell = 0
+            recv_data_time = time.time()
+            logger.debug(f"self.sequence_number {self.sequence_number} "
+                         f"recv_data_time {recv_data_time - self.start_time}")
+
+            if self.obj_body == self.ACK_STOP:
+                # Indicates the current TCP link receives a STOP signal and remove from the transport_dict
+                _transport = self.factory.transport_dict.pop(self.transport, None)
+                logger.debug(f"接收到b'STOP_' self.sequence_number {self.sequence_number} ")
+            self.post_process()
+
+
+class MessageServerFactory(protocol.ServerFactory):
+    def __init__(self) -> None:
+        """
+        transport_dict: links that have not completed data transmission.
+        transport_list: Records all TCP links. Appends TCP link to the transport list when a new TCP link is established.
+        """
+        self.transport_dict = {}
+        self.transport_list = []
+
+    def is_all_connection_closed(self):
+        return len(self.transport_dict) == 0

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/ssl_config.py

@@ -1,10 +1,10 @@
-cipher_list = ":".join([
-    'ECDHE-ECDSA-AES128-GCM-SHA256',
-    'ECDHE-RSA-AES128-GCM-SHA256',
-    'ECDHE-ECDSA-AES256-GCM-SHA384',
-    'ECDHE-RSA-AES256-GCM-SHA384',
-    'ECDHE-ECDSA-CHACHA20-POLY1305',
-    'ECDHE-RSA-CHACHA20-POLY1305',
-    'DHE-RSA-AES128-GCM-SHA256',
-    'DHE-RSA-AES256-GCM-SHA384'
-]).encode()
+cipher_list = ":".join([
+    'ECDHE-ECDSA-AES128-GCM-SHA256',
+    'ECDHE-RSA-AES128-GCM-SHA256',
+    'ECDHE-ECDSA-AES256-GCM-SHA384',
+    'ECDHE-RSA-AES256-GCM-SHA384',
+    'ECDHE-ECDSA-CHACHA20-POLY1305',
+    'ECDHE-RSA-CHACHA20-POLY1305',
+    'DHE-RSA-AES128-GCM-SHA256',
+    'DHE-RSA-AES256-GCM-SHA384'
+]).encode()

msprobe/pytorch/bench_functions/__init__.py

@@ -1,15 +1,15 @@
-import os
-from pkgutil import iter_modules
-from importlib import import_module
-
-"""
-gpu and cpu not implement benchmark function, supplementary benchmarking function implementation
-"""
-
-package_path = os.path.dirname(os.path.realpath(__file__))
-for _, module_name, _ in iter_modules([package_path]):
-    module = import_module(f"{__name__}.{module_name}")
-    for attr_name in dir(module):
-        attr = getattr(module, attr_name)
-        if callable(attr) and "npu_custom" not in attr_name:
-            globals()[attr_name] = attr
+import os
+from pkgutil import iter_modules
+from importlib import import_module
+
+"""
+gpu and cpu not implement benchmark function, supplementary benchmarking function implementation
+"""
+
+package_path = os.path.dirname(os.path.realpath(__file__))
+for _, module_name, _ in iter_modules([package_path]):
+    module = import_module(f"{__name__}.{module_name}")
+    for attr_name in dir(module):
+        attr = getattr(module, attr_name)
+        if callable(attr) and "npu_custom" not in attr_name:
+            globals()[attr_name] = attr

msprobe/pytorch/bench_functions/apply_adam_w.py

@@ -1,28 +1,28 @@
-import torch
-
-
-def npu_apply_adam_w(beta1_power, beta2_power, lr, weight_decay,
-                     beta1, beta2, eps, grad, max_grad_norm, amsgrad, maximize, out):
-    var, m, v = out
-    if amsgrad:
-        max_grad_norm = (torch.rand(var.shape) * 10.0 - 5.0).to(var.dtype)
-    beta1_power_out = beta1_power * beta1
-    beta2_power_out = beta2_power * beta2
-    var_t = var * (1 + (-lr * weight_decay))
-    gt = -grad if maximize else grad
-    m_out = m * beta1 - (beta1 + (-1)) * gt
-    v_out = v * beta2 - (beta2 + (-1)) * gt * gt
-
-    if amsgrad:
-        max_grad_norm_out = torch.max(max_grad_norm, v_out)
-        if (1 - beta2_power_out) == 0:
-            beta2_power_out -= eps
-        denom = torch.sqrt(torch.div(max_grad_norm_out, (1 - beta2_power_out))) + eps
-    else:
-        vraintain = torch.div(v_out, (1 - beta2_power_out))
-        denom = torch.sqrt(vraintain) + eps
-
-    if (1 - beta1_power_out) == 0:
-        beta1_power_out -= eps
-    var_out = var_t + torch.div(-lr * m_out, (1 - beta1_power_out)).div(denom)
-    return var_out.cpu(), m_out.cpu(), v_out.cpu()
+import torch
+
+
+def npu_apply_adam_w(beta1_power, beta2_power, lr, weight_decay,
+                     beta1, beta2, eps, grad, max_grad_norm, amsgrad, maximize, out):
+    var, m, v = out
+    if amsgrad:
+        max_grad_norm = (torch.rand(var.shape) * 10.0 - 5.0).to(var.dtype)
+    beta1_power_out = beta1_power * beta1
+    beta2_power_out = beta2_power * beta2
+    var_t = var * (1 + (-lr * weight_decay))
+    gt = -grad if maximize else grad
+    m_out = m * beta1 - (beta1 + (-1)) * gt
+    v_out = v * beta2 - (beta2 + (-1)) * gt * gt
+
+    if amsgrad:
+        max_grad_norm_out = torch.max(max_grad_norm, v_out)
+        if (1 - beta2_power_out) == 0:
+            beta2_power_out -= eps
+        denom = torch.sqrt(torch.div(max_grad_norm_out, (1 - beta2_power_out))) + eps
+    else:
+        vraintain = torch.div(v_out, (1 - beta2_power_out))
+        denom = torch.sqrt(vraintain) + eps
+
+    if (1 - beta1_power_out) == 0:
+        beta1_power_out -= eps
+    var_out = var_t + torch.div(-lr * m_out, (1 - beta1_power_out)).div(denom)
+    return var_out, m_out, v_out

msprobe/pytorch/bench_functions/confusion_transpose.py

@@ -1,19 +1,19 @@
-def npu_confusion_transpose(data, perm, shape, transpose_first):
-    if transpose_first:
-        output = data.permute(*perm).contiguous().view(shape)
-    else:
-        output = data.view(shape).permute(*perm)
-    return output.cpu()
-
-
-def npu_confusion_transpose_backward(grad, perm, shape, transpose_first):
-    shape_cal = shape if transpose_first else [shape[perm_dim] for perm_dim in perm]
-    perm_cal = [0] * len(perm)
-    for i, perm_dim in enumerate(perm):
-        perm_cal[perm_dim] = i
-
-    if transpose_first:
-        result = grad.permute(*perm_cal).reshape(shape_cal)
-    else:
-        result = grad.reshape(shape_cal).permute(*perm_cal)
-    return result.cpu()
+def npu_confusion_transpose(data, perm, shape, transpose_first):
+    if transpose_first:
+        output = data.permute(*perm).contiguous().view(shape)
+    else:
+        output = data.view(shape).permute(*perm)
+    return output
+
+
+def npu_confusion_transpose_backward(grad, perm, shape, transpose_first):
+    shape_cal = shape if transpose_first else [shape[perm_dim] for perm_dim in perm]
+    perm_cal = [0] * len(perm)
+    for i, perm_dim in enumerate(perm):
+        perm_cal[perm_dim] = i
+
+    if transpose_first:
+        result = grad.permute(*perm_cal).reshape(shape_cal)
+    else:
+        result = grad.reshape(shape_cal).permute(*perm_cal)
+    return result.cpu()