mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.4.0__cp39-cp39-win_amd64.whl

mindspore/_extends/remote/kernel_build_server.py

@@ -154,6 +154,8 @@ class AkgBuilderBase():
             messager.send_ack()
             wait_time_str = messager.get_message()
             messager.send_ack()
+            self.platform = messager.get_message()
+            messager.send_ack()
             self.create(int(process_num_str), int(wait_time_str))
         elif arg == 'AKG/ATTR':
             messager.send_ack()

mindspore/_profiler.py

@@ -0,0 +1,30 @@
+# Copyright 2023 Huawei Technologies Co., Ltd
+#
+# 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.
+# ============================================================================
+"""python profiler"""
+
+
+from mindspore._c_expression import PythonProfilerRecorder
+
+
+def ms_profile(fn):
+    def wrap_profile_function(*args, **kwargs):
+        recorder = PythonProfilerRecorder(fn.__qualname__)
+
+        recorder.record_start()
+        ret = fn(*args, **kwargs)
+        recorder.record_end()
+        return ret
+
+    return wrap_profile_function

mindspore/amp.py

@@ -20,8 +20,11 @@ from abc import ABC, abstractmethod
 from mindspore.common import mutable
 from mindspore.ops._primitive_cache import _get_cache_prim
 from mindspore.ops.operations.math_ops import NPUGetFloatStatusV2, NPUClearFloatStatusV2
+from mindspore.ops.operations.nn_ops import AllFinite
+from mindspore.run_check._check_version import AscendEnvChecker
 from mindspore import _checkparam as validator
 from mindspore._c_expression import MSContext
+from mindspore import log as logger
 from .common import dtype as mstype
 from . import context
 from . import ops
@@ -30,22 +33,26 @@ from .common.api import jit_class, jit
 from .common.parameter import Parameter
 from .common.tensor import Tensor
 from .train.loss_scale_manager import DynamicLossScaleManager, LossScaleManager, FixedLossScaleManager
-from .train.amp import build_train_network, auto_mixed_precision, custom_mixed_precision,\
+from .train.amp import build_train_network, auto_mixed_precision, custom_mixed_precision, \
     get_white_list, get_black_list
 
-
 _hypermap = ops.HyperMap()
 _partial = ops.Partial()
 
 
 @constexpr
-def _ascend_910A_target():
+def _ascend_target():
+    return context.get_context("device_target") == "Ascend"
+
+
+@constexpr
+def _ascend_910a_target():
     return MSContext.get_instance().get_ascend_soc_version() == "ascend910"
 
 
 @constexpr
-def _ascend_910B_target():
-    return MSContext.get_instance().get_ascend_soc_version() == "ascend910b"
+def _ascend_910b_target():
+    return MSContext.get_instance().get_ascend_soc_version() in ["ascend910b", "ascend910_93"]
 
 
 @constexpr
@@ -53,6 +60,31 @@ def _gpu_target():
     return context.get_context("device_target") == "GPU"
 
 
+@constexpr
+def _enable_all_finite():
+    """check whether enable all finite"""
+    logger.debug("Start enable all finite.")
+    if _ascend_target():
+        checker = AscendEnvChecker(None)
+        if not checker.check_custom_version():
+            logger.debug("Disable AllFinite due to version check failure.")
+            return False
+    runtime_conf = os.environ.get('MS_DEV_RUNTIME_CONF')
+    global_jit_config = context.get_jit_config()
+    if runtime_conf is not None and ("all_finite:True" in runtime_conf or "all_finite:true" in runtime_conf):
+        logger.debug("Enable AllFinite through the environment variable MS_DEV_RUNTIME_CONF.")
+        return True
+
+    if runtime_conf is not None and ("all_finite:False" in runtime_conf or "all_finite:false" in runtime_conf):
+        logger.debug("Disable AllFinite through the environment variable MS_DEV_RUNTIME_CONF.")
+        return False
+
+    if global_jit_config:
+        logger.debug("Current global jit config is: {}".format(global_jit_config["jit_level"]))
+        return global_jit_config["jit_level"] == "O0" or global_jit_config["jit_level"] == "O1"
+    return False
+
+
 def _grad_unscale(scale, grad):
     return grad * ops.Reciprocal()(scale).astype(grad.dtype)
 
@@ -79,22 +111,27 @@ def _overflow(inputs):
 
 
 @jit
-def _all_finite(inputs, check_overflow_mode):
+def _all_finite(inputs, check_overflow_mode, enable_allfinite):
     """all finite check"""
-    if (_ascend_910A_target()) or \
-       (_ascend_910B_target() and check_overflow_mode != "INFNAN_MODE"):
-        status = Tensor([0] * 8, mstype.int32)
-        status = ops.depend(status, inputs)
-        get_status = _get_cache_prim(NPUGetFloatStatusV2)()(status)
-        status = ops.depend(status, get_status)
-        clear_status = _get_cache_prim(NPUClearFloatStatusV2)()(status)
-        get_status = ops.depend(get_status, clear_status)
-        status_finite = get_status.equal(Tensor(0, mstype.int32)).all()
-        return status_finite
-
-    outputs = _hypermap(_partial(_overflow), inputs)
-    flag_sum = ops.addn(outputs).reshape(())
-    status_finite = ops.less(flag_sum, 1)
+    if _ascend_target():
+        if (_ascend_910a_target()) or \
+                (_ascend_910b_target() and check_overflow_mode == "SATURATION_MODE"):
+            status = Tensor([0] * 8, mstype.int32)
+            status = ops.depend(status, inputs)
+            get_status = _get_cache_prim(NPUGetFloatStatusV2)()(status)
+            status = ops.depend(status, get_status)
+            clear_status = _get_cache_prim(NPUClearFloatStatusV2)()(status)
+            get_status = ops.depend(get_status, clear_status)
+            status_finite = get_status.equal(Tensor(0, mstype.int32)).all()
+            return status_finite
+
+    status_finite = False
+    if enable_allfinite:
+        status_finite = ~AllFinite()(inputs)  # pylint: disable=invalid-unary-operand-type
+    else:
+        outputs = _hypermap(_partial(_overflow), inputs)
+        flag_sum = ops.addn(outputs).reshape(())
+        status_finite = ops.less(flag_sum, 1)
     return status_finite
 
 
@@ -126,11 +163,11 @@ def all_finite(inputs):
 
     Tutorial Examples:
         - `Automatic Mix Precision - Loss Scaling
-          <https://mindspore.cn/tutorials/en/r2.2/advanced/mixed_precision.html#loss-scaling>`_
+          <https://mindspore.cn/tutorials/en/master/beginner/mixed_precision.html#loss-scaling>`_
     """
     inputs = mutable(inputs)
     _check_overflow_mode = os.environ.get('MS_ASCEND_CHECK_OVERFLOW_MODE')
-    return _all_finite(inputs, _check_overflow_mode)
+    return _all_finite(inputs, _check_overflow_mode, _enable_all_finite())
 
 
 @jit_class
@@ -142,7 +179,7 @@ class LossScaler(ABC):
     to scale and unscale the loss value and gradients to avoid overflow, `adjust` is used to update the
     loss scale value.
 
-    For more information, refer to the `tutorials  <https://mindspore.cn/tutorials/en/r2.2/advanced/
+    For more information, refer to the `tutorials  <https://mindspore.cn/tutorials/en/master/beginner/
     mixed_precision.html#loss-scaling>`_.
 
     .. warning::
@@ -173,6 +210,7 @@ class LossScaler(ABC):
         >>>
         >>> loss_scaler = MyLossScaler(1024)
     """
+
     @abstractmethod
     def scale(self, inputs):
         """
@@ -235,6 +273,7 @@ class StaticLossScaler(LossScaler):
         (Tensor(shape=[2], dtype=Float16, value= [ 1.4648e-03,  9.7656e-04]),
         Tensor(shape=[1], dtype=Float16, value= [ 1.1721e-03]))
     """
+
     def __init__(self, scale_value):
         scale_value = validator.check_value_type("scale_value", scale_value, [float, int])
         if scale_value < 1.0:
@@ -313,6 +352,7 @@ class DynamicLossScaler(LossScaler):
         >>> print(loss_scaler.scale_value.asnumpy())
         512.0
     """
+
     def __init__(self, scale_value, scale_factor, scale_window):
         scale_value = validator.check_value_type("scale_value", scale_value, [float, int])
         if scale_value < 1.0:
@@ -334,7 +374,7 @@ class DynamicLossScaler(LossScaler):
 
         Tutorial Examples:
             - `Automatic Mix Precision - Loss Scaling
-              <https://mindspore.cn/tutorials/en/r2.2/advanced/mixed_precision.html#loss-scaling>`_
+              <https://mindspore.cn/tutorials/en/master/beginner/mixed_precision.html#loss-scaling>`_
         """
         inputs = mutable(inputs)
         return _grad_scale_map(self.scale_value, inputs)
@@ -351,7 +391,7 @@ class DynamicLossScaler(LossScaler):
 
         Tutorial Examples:
             - `Automatic Mix Precision - Loss Scaling
-              <https://mindspore.cn/tutorials/en/r2.2/advanced/mixed_precision.html#loss-scaling>`_
+              <https://mindspore.cn/tutorials/en/master/beginner/mixed_precision.html#loss-scaling>`_
         """
         inputs = mutable(inputs)
         return _grad_unscale_map(self.scale_value, inputs)
@@ -365,7 +405,7 @@ class DynamicLossScaler(LossScaler):
 
         Tutorial Examples:
             - `Automatic Mix Precision - Loss Scaling
-              <https://mindspore.cn/tutorials/en/r2.2/advanced/mixed_precision.html#loss-scaling>`_
+              <https://mindspore.cn/tutorials/en/master/beginner/mixed_precision.html#loss-scaling>`_
         """
         one = ops.ones((), self.scale_value.dtype)
         scale_mul_factor = self.scale_value * self.scale_factor
@@ -384,6 +424,7 @@ class DynamicLossScaler(LossScaler):
         ops.assign(self.counter, counter)
         return True
 
+
 __all__ = [
     "DynamicLossScaleManager", "LossScaleManager", "FixedLossScaleManager",
     "build_train_network", "DynamicLossScaler", "StaticLossScaler", "LossScaler",

mindspore/boost/adasum.py

@@ -23,7 +23,7 @@ from mindspore.communication.management import create_group
 from mindspore.ops import composite as C
 from mindspore.ops import functional as F
 from mindspore.ops import operations as P
-from mindspore.ops.operations._inner_ops import Send, Receive
+from mindspore.ops import Send, Receive
 
 
 __all__ = ["AdaSum"]

mindspore/boost/base.py

@@ -183,7 +183,7 @@ class ParameterProcess:
     Examples:
         >>> import numpy as np
         >>> from mindspore import Tensor, Parameter, nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore.boost import ParameterProcess
         >>>
         >>> class Net(nn.Cell):

mindspore/boost/boost_cell_wrapper.py

@@ -136,7 +136,7 @@ class BoostTrainOneStepCell(TrainOneStepCell):
         >>> from mindspore import boost
         >>> from mindspore import nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits()
         >>> optim = nn.Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
@@ -435,7 +435,7 @@ class BoostTrainOneStepWithLossScaleCell(BoostTrainOneStepCell):
     Examples:
         >>> import numpy as np
         >>> from mindspore import Tensor, Parameter, nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore.nn import WithLossCell
         >>> from mindspore import dtype as mstype
         >>> from mindspore import boost

mindspore/boost/grad_freeze.py

@@ -206,7 +206,7 @@ class GradientFreeze:
     Examples:
         >>> import numpy as np
         >>> from mindspore import Tensor, Parameter, nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore.nn import WithLossCell
         >>> from mindspore import dtype as mstype
         >>> from mindspore import boost
@@ -347,7 +347,7 @@ def freeze_cell(reducer_flag, network, optimizer, sens, grad, use_grad_accumulat
     Examples:
         >>> import numpy as np
         >>> from mindspore import Tensor, Parameter, nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore.boost.grad_freeze import freeze_cell, FreezeOpt
         >>>
         >>> class Net(nn.Cell):

mindspore/boost/group_loss_scale_manager.py

@@ -94,7 +94,7 @@ class GroupLossScaleManager(Cell):
         ...                        loss_scale_manager=loss_scale_manager,
         ...                        boost_level="O1", boost_config_dict=boost_config_dict)
         >>> # Create the dataset taking MNIST as an example. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/mnist.py
         >>> dataset = create_dataset()
         >>> model.train(2, dataset)
     """

mindspore/boost/less_batch_normalization.py

@@ -96,7 +96,7 @@ class LessBN(Cell):
     Examples:
         >>> import numpy as np
         >>> from mindspore import Tensor, Parameter, nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> from mindspore.nn import WithLossCell
         >>> from mindspore import dtype as mstype
         >>> from mindspore import boost
@@ -125,11 +125,14 @@ class LessBN(Cell):
     def __init__(self, network, fn_flag=False):
         super(LessBN, self).__init__()
         self.network = network
-        self.network.set_boost("less_bn")
-        self.network.update_cell_prefix()
-        if fn_flag:
-            self._convert_to_less_bn_net(self.network)
-        self.network.add_flags(defer_inline=True)
+        if "less_bn" not in self.network.get_flags():
+            self.network.set_boost("less_bn")
+            self.network.update_cell_prefix()
+            if fn_flag:
+                self._convert_to_less_bn_net(self.network)
+            # Please use this interface with caution as it may cause uncontrollable problem.
+            # The purpose here is to prevent the "less_bn" flag from disappearing.
+            self.network.add_flags(defer_inline=True)
 
     @staticmethod
     def _convert_dense(subcell):

mindspore/common/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2021 Huawei Technologies Co., Ltd
+# Copyright 2020-2024 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,13 +15,13 @@
 """Top-level reference to dtype of common module."""
 from __future__ import absolute_import
 from mindspore.common import dtype
-from mindspore.common.api import ms_function, ms_memory_recycle, ms_class, jit, jit_class, _no_grad
+from mindspore.common.api import ms_function, ms_memory_recycle, ms_class, jit, jit_class, _no_grad, flops_collection
 from mindspore.common.dtype import Type, int8, byte, int16, short, int32, intc, int64, intp, \
     uint8, ubyte, uint16, ushort, uint32, uintc, uint64, uintp, float16, half, \
     float32, single, float64, bfloat16, double, bool_, float_, list_, tuple_, int_, \
     uint, number, tensor_type, string, type_none, TensorType, Int, \
     complex64, complex128, dtype_to_nptype, _null, _NullType, \
-    dtype_to_pytype, pytype_to_dtype, get_py_obj_dtype, QuantDtype
+    dtype_to_pytype, pytype_to_dtype, get_py_obj_dtype, QuantDtype, qint4x2
 from mindspore.common.dump import set_dump
 from mindspore.common.parameter import Parameter, ParameterTuple
 from mindspore.common.seed import set_seed, get_seed
@@ -30,7 +30,14 @@ from mindspore.common.sparse_tensor import RowTensor, RowTensorInner, SparseTens
 from mindspore.common.mutable import mutable
 from mindspore.common.jit_config import JitConfig
 from mindspore.common.lazy_inline import lazy_inline
+from mindspore.common.no_inline import no_inline
 from mindspore.common.mindir_util import load_mindir, save_mindir
+from mindspore.common.symbol import Symbol
+from mindspore.common.recompute import recompute
+from mindspore.common import generator
+from mindspore.common.generator import (
+    Generator, default_generator, seed, manual_seed, initial_seed, get_rng_state, set_rng_state)
+from mindspore.ops.function.array_func import is_tensor, from_numpy
 
 # symbols from dtype
 __all__ = [
@@ -57,17 +64,23 @@ __all__ = [
     # __method__ from dtype
     "dtype_to_nptype", "dtype_to_pytype",
     "pytype_to_dtype", "get_py_obj_dtype",
-    "bfloat16",
+    "bfloat16", "qint4x2"
 ]
 
 __all__.extend([
-    "tensor", "Tensor", "RowTensor", "SparseTensor", "COOTensor", "CSRTensor", # tensor
+    "tensor", "Tensor", "RowTensor", "SparseTensor", "COOTensor", "CSRTensor",  # tensor
     "ms_function", "ms_class", 'jit', 'jit_class', '_no_grad',  # api
     "Parameter", "ParameterTuple",  # parameter
     "dtype",
-    "set_seed", "get_seed",  # random seed
+    "set_seed", "get_seed", "manual_seed", # random seed
     "set_dump",
     "ms_memory_recycle",
     "mutable", "JitConfig",
-    "lazy_inline", "load_mindir", "save_mindir"
+    "flops_collection",
+    "lazy_inline", "load_mindir", "save_mindir",
+    "no_inline",
+    "Symbol",
+    "recompute",
+    "is_tensor", "from_numpy",
 ])
+__all__.extend(generator.__all__)

mindspore/common/_jit_fallback_utils.py

@@ -16,7 +16,6 @@
 # ============================================================================
 """common _jit_fallback_utils."""
 
-
 _local_value_nodes = {}
 
 
@@ -81,8 +80,8 @@ def list_inplace_pop(list_obj, index):
     # This will be removed after empty list problem is solved.
     if isinstance(list_obj, tuple) and not list_obj == 0:
         list_obj = []
-    list_obj.pop(index)
-    return list_obj
+    ret = list_obj.pop(index)
+    return list_obj, ret
 
 
 def list_inplace_reverse(list_obj):

mindspore/common/_pijit_context.py

@@ -0,0 +1,190 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# 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.
+# ============================================================================
+"""Define pijit context"""
+
+import inspect
+import types
+import functools
+import importlib.util
+import mindspore
+from mindspore import log as logger
+from mindspore.common.jit_config import JitConfig
+from mindspore._c_expression import GraphExecutor_, jit_mode_pi_enable, jit_mode_pi_disable, pi_jit_set_context
+
+
+def _update_graph_executor_config(jit_config):
+    """Update GraphExecutor jit_config"""
+    if isinstance(jit_config, JitConfig):
+        jit_config = jit_config.jit_config_dict
+    if not isinstance(jit_config, dict):
+        return
+    valid_config = {}
+    for k, v in jit_config.items():
+        valid_config[str(k)] = str(v)
+    GraphExecutor_.get_instance().set_jit_config(JitConfig(**valid_config).jit_config_dict)
+
+
+class PIJitCaptureContext:
+    """
+    Context manager for pijit graph capture
+    """
+
+    def __init__(self, jit_config=None, input_signature=None):
+        _update_graph_executor_config(jit_config)
+        config = {}
+        if isinstance(jit_config, JitConfig):
+            config.update(jit_config.jit_config_dict)
+        elif jit_config is not None:
+            config.update(jit_config)
+
+        self.config = config
+        self.input_signature = input_signature
+        self.ret = None
+        self.fn = None
+        self._init_arg = iter([self.config, self.input_signature])
+
+        if not SKIP_RULES:
+            return
+        pi_jit_set_context(wrapper=self._wrapper(),
+                           skip_files=_get_skip_files(),
+                           skip_codes=SKIP_RULES["codes"])
+        SKIP_RULES.clear()
+
+    @staticmethod
+    def _is_unsupported(fn):
+        # generator, coroutine, awaitable and a function that return them is unsupported
+        return inspect.isgeneratorfunction(fn) or inspect.iscoroutinefunction(fn) \
+            or inspect.isasyncgenfunction(fn) or inspect.isawaitable(fn)
+
+    def _wrapper(self):
+        def _fn(*args, **kwds):
+            with self:
+                self.ret = self.fn(*args, **kwds)
+                return self.ret
+        return _fn
+
+    def __call__(self, fn):
+        if isinstance(fn, type) and issubclass(fn, mindspore.nn.Cell):
+            fn.construct = self(fn.construct)
+            return fn
+        if isinstance(fn, mindspore.nn.Cell):
+            type(fn).construct = self(type(fn).construct)
+            return fn
+        if isinstance(fn, types.MethodType):
+            return types.MethodType(self(fn.__func__), fn.__self__)
+        if not isinstance(fn, types.FunctionType) or self._is_unsupported(fn):
+            logger.warning("unsupported function type" + str(fn))
+            return fn
+
+        try:
+            if inspect.getmodule(fn.__code__).__name__.startswith("mindspore"):
+                return fn
+        finally:
+            pass
+
+        _fn = self._wrapper()
+        if fn.__code__ is _fn.__code__:
+            fn = fn.__closure__[0].cell_contents.fn
+        self.fn = fn
+        return functools.wraps(fn)(_fn)
+
+    def __enter__(self):
+        pi_jit_set_context(self.fn, *self._init_arg)
+        jit_mode_pi_enable()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pi_jit_set_context(None)
+        jit_mode_pi_disable()
+
+
+def _get_skip_files():
+    """
+    Get skip files by SKIP_RULES
+    """
+    def _filter(path: str):
+        if path.endswith("__init__.py"):
+            return path[0:-11]
+        return path
+
+    # not import these modules, only find it
+    find = importlib.util.find_spec
+
+    files = [*SKIP_RULES["skip_dirs"]]
+    files += [_filter(find(m).origin) for m in SKIP_RULES["builtins"]]
+    for i in SKIP_RULES["third_party"]:
+        spec = find(i)
+        if spec is None:
+            continue
+        files.append(_filter(spec.origin))
+
+    return tuple(files)
+
+
+# complete the skip list...
+SKIP_RULES = {
+    "skip_dirs": (
+        "<frozen importlib",
+        "<__array_function__ internals>",
+        "<string>",
+    ),
+    "builtins": (
+        "mindspore",  # not capture any function of mindspore unless it's called by user
+        "abc",
+        "ast",
+        "codecs",
+        "collections",
+        "contextlib",
+        "copy",
+        "copyreg",
+        "dataclasses",
+        "enum",
+        "functools",
+        "glob",
+        "importlib",
+        "inspect",
+        "linecache",
+        "logging",
+        "multiprocessing",
+        "operator",
+        "os",
+        "posixpath",
+        "random",
+        "re",
+        "selectors",
+        "signal",
+        "tempfile",
+        "threading",
+        "tokenize",
+        "traceback",
+        "types",
+        "typing",
+        "unittest",
+        "weakref",
+        "_collections_abc",
+        "_weakrefset",
+        # others...
+        "sre_compile",
+        "sre_parse",
+        "genericpath",
+    ),
+    "third_party": (
+        "numpy",
+        "pandas",
+        "sklearn",
+        "tqdm",
+        "tree",
+    ),
+    "codes": (),
+}

mindspore/common/_register_for_adapter.py

@@ -24,6 +24,7 @@ class Registry:
         self._tensor = None
         self._parameter = None
         self._convert_map = {}
+        self.convert_adapter_tensor_map = {}
 
     @property
     def tensor(self):
@@ -63,5 +64,11 @@ class Registry:
             raise ValueError(f"Expect a dict type, but got {type(value)}.")
         self._convert_map = value
 
+    def register_convert_adapter_tensor_map(self, value):
+        """Register the convert map for adapter tensor."""
+        if not isinstance(value, dict):
+            raise ValueError(f"Expect a dict type, but got {type(value)}.")
+        self.convert_adapter_tensor_map = value
+
 
 ms_adapter_registry = Registry()