mindspore 2.4.10__cp311-cp311-win_amd64.whl → 2.6.0rc1__cp311-cp311-win_amd64.whl

mindspore/common/_pijit_context.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ============================================================================
 """Define pijit context"""
-
+# pylint: disable=missing-docstring
 import inspect
 import types
 import functools
@@ -21,6 +21,7 @@ import importlib.util
 import mindspore
 from mindspore import log as logger
 from mindspore.common.jit_config import JitConfig
+from mindspore._c_expression import PreJit
 from mindspore._c_expression import GraphExecutor_, jit_mode_pi_enable, jit_mode_pi_disable, pi_jit_set_context
 
 
@@ -49,6 +50,10 @@ class PIJitCaptureContext:
         elif jit_config is not None:
             config.update(jit_config)
 
+        disable_pijit = config.get('_disable_pijit', None)
+        if disable_pijit is not None and not callable(disable_pijit):
+            raise TypeError(f"The config '_disable_pijit' must be callable but got {disable_pijit}")
+
         self.config = config
         self.input_signature = input_signature
         self.ret = None
@@ -70,6 +75,10 @@ class PIJitCaptureContext:
 
     def _wrapper(self):
         def _fn(*args, **kwds):
+            PreJit(args, kwds)
+            disable_pijit = self.config.get('_disable_pijit', None)
+            if disable_pijit is not None and disable_pijit(args, kwds):
+                return self.fn(*args, **kwds)
             with self:
                 self.ret = self.fn(*args, **kwds)
                 return self.ret
@@ -88,11 +97,11 @@ class PIJitCaptureContext:
             logger.warning("unsupported function type" + str(fn))
             return fn
 
-        try:
-            if inspect.getmodule(fn.__code__).__name__.startswith("mindspore"):
+        module = inspect.getmodule(fn.__code__)
+        if module is not None and module.__name__.startswith("mindspore"):
+            if fn.__code__.co_name != 'after_grad':
+                # Use PIJit for mindspore api, please use PSJit
                 return fn
-        finally:
-            pass
 
         _fn = self._wrapper()
         if fn.__code__ is _fn.__code__:

mindspore/common/_register_for_tensor.py

@@ -16,7 +16,7 @@
 """Registry the relation."""
 
 from __future__ import absolute_import
-from mindspore._c_expression import Tensor as Tensor_
+from mindspore._c_expression import TensorPy as Tensor_
 
 
 class Registry:
@@ -43,4 +43,3 @@ class Registry:
 
 
 tensor_operator_registry = Registry()
-tensor_operator_registry_for_mint = Registry()

mindspore/common/_stub_tensor.py

@@ -18,9 +18,11 @@ import inspect
 from functools import reduce
 from mindspore.common.tensor import Tensor
 from mindspore.common.dtype import type_size_in_bytes
+import mindspore._c_expression as c_expression
 from mindspore._c_expression import TensorNode, SequenceNode, NoneTypeNode, AnyTypeNode
-from mindspore._c_expression import Tensor as Tensor_
+from mindspore._c_expression import TensorPy as Tensor_
 from mindspore.common.api import _convert_python_data
+from mindspore.common._tensor_cpp_method import tensor_cpp_methods
 
 
 def _stub_member(var, init):
@@ -56,16 +58,10 @@ class StubTensor:
     init = _stub_member("init", None)
     init_finished = _stub_member("init_finished", False)
     virtual_flag = _stub_member("virtual_flag", False)
-    adapter_flag = _stub_member("adapter_flag", False)
     parent_tensor_ = _stub_member("parent_tensor_", None)
     index_of_parent_ = _stub_member("index_of_parent_", None)
     slice_num_of_persistent_data_ = _stub_member("slice_num_of_persistent_data_", None)
     slice_shape_of_persistent_data_ = _stub_member("slice_shape_of_persistent_data_", None)
-    # auto gradient information
-    _grad = _stub_member("_grad", None)
-    _grad_fn = _stub_member("_grad_fn", None)
-    _requires_grad = _stub_member("_requires_grad", False)
-    _retain_grad = _stub_member("_retain_grad", False)
 
     def __init__(self, stub=None, tensor=None):
         self.stub = stub
@@ -143,21 +139,40 @@ class StubTensor:
         """
         return self.ndim
 
+
+    def is_contiguous(self):
+        if self.stub:
+            return self.stub.get_value().is_contiguous()
+        return self.tensor.is_contiguous()
+
+    def set_cast_dtype(self):
+        if self.stub:
+            return self.stub.get_value().set_cast_dtype()
+        return self.tensor.set_cast_dtype()
+
+    def storage_offset(self):
+        if self.stub:
+            return self.stub.get_value().storage_offset()
+        return self.tensor.storage_offset()
+
+    def _need_contiguous(self):
+        if self.stub:
+            return self.stub.get_value()._need_contiguous()  # pylint: disable=protected-access
+        return self.tensor._need_contiguous()  # pylint: disable=protected-access
+
     asnumpy = _stub_method(Tensor.asnumpy)
     is_persistent_data = _stub_method(Tensor.is_persistent_data)
     asnumpy_of_slice_persistent_data = _stub_method(Tensor.asnumpy_of_slice_persistent_data)
     slice_num_of_persistent_data = _stub_method(Tensor.slice_num_of_persistent_data)
     slice_shape_of_persistent_data = _stub_method(Tensor.slice_shape_of_persistent_data)
     flush_from_cache = _stub_method(Tensor.flush_from_cache)
-    contiguous = _stub_method(Tensor.contiguous)
-    is_contiguous = _stub_method(Tensor.is_contiguous)
     register_hook = _stub_method(Tensor.register_hook)
 
     def stub_sync(self):
         """sync real tensor."""
         if self.stub:
             val = self.stub.get_value()
-            self.tensor = Tensor(val, internal=True)
+            self.tensor = Tensor(val)
             if hasattr(self, "member_cache"):
                 for k, v in self.member_cache.items():
                     setattr(self.tensor, k, v)
@@ -173,8 +188,9 @@ class StubTensor:
     def __setstate__(self, state):
         value = state.pop("value")
         self.stub = None
-        self.tensor = Tensor(value, internal=True)
+        self.tensor = Tensor(value)
 
+no_stub_sync_cpp_api = ["set_cast_dtype", "storage_offset", "is_contiguous", "_need_contiguous"]
 
 def _init_stub_tensor_api():
     """adapt to python tensor and cpp tensor api"""
@@ -182,15 +198,15 @@ def _init_stub_tensor_api():
     cpp_tensor_func = dir(Tensor_)
     for attr in need_init_func:
         func = inspect.getattr_static(Tensor, attr)
-        if attr in cpp_tensor_func:
+        if attr in cpp_tensor_func and attr not in tensor_cpp_methods:
             # for cpp tensor api, we always need to sync for real tensor first
             setattr(StubTensor, attr, _stub_method(func))
-        else:
+        elif attr not in no_stub_sync_cpp_api:
             setattr(StubTensor, attr, func)
 
 
 _init_stub_tensor_api()
-
+c_expression.addStubTensorMethods()
 
 def _convert_stub(stub):
     "convert stub to StubNode or Value"

mindspore/common/_tensor_cpp_method.py

@@ -0,0 +1,17 @@
+# 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.
+# ============================================================================
+"""Add tensor cpp methods for stub tensor"""
+
+tensor_cpp_methods = ['abs', 'absolute', '__abs__', 'acos', 'arccos', 'acosh', 'arccosh', 'add', '__add__', 'addbmm', 'addcdiv', 'addmm', 'addmv', 'add_', '__iadd__', 'all', 'allclose', 'any', 'argmax', 'argmin', 'argsort', 'asin', 'arcsin', 'asinh', 'arcsinh', 'atan', 'arctan', 'atan2', 'arctan2', 'atanh', 'arctanh', 'baddbmm', 'bincount', 'bitwise_and', '__and__', 'bitwise_not', 'bitwise_or', '__or__', 'bitwise_xor', '__xor__', 'ceil', 'chunk', 'clamp', 'clip', 'clone', 'copy_', 'cos', 'cosh', 'count_nonzero', 'cumsum', 'diag', 'div', 'divide', 'div_', '__itruediv__', 'dot', 'eq', 'erf', 'erfc', 'exp', 'expand_as', 'expm1', 'exp_', 'fill_', 'fill_diagonal_', 'flatten', 'floor', 'floor_divide', 'floor_divide_', '__ifloordiv__', 'fmod', 'frac', 'gather', 'gcd', 'greater', 'gt', 'greater_equal', 'ge', 'hardshrink', 'histc', 'index_add', 'index_select', 'inverse', 'isclose', 'isfinite', 'isinf', 'isneginf', 'kthvalue', 'lerp', 'less', 'lt', 'less_equal', 'le', 'log', 'log10', 'log1p', 'log2', 'logaddexp', 'logaddexp2', 'logical_and', 'logical_not', 'logical_or', 'logical_xor', 'logsumexp', 'log_', 'masked_fill', 'masked_fill_', 'masked_select', 'matmul', 'max', 'maximum', 'mean', 'median', 'min', 'minimum', 'mm', 'mul', 'mul_', '__imul__', 'nansum', 'nan_to_num', 'narrow', 'neg', 'negative', 'new_ones', 'new_zeros', 'not_equal', 'ne', 'outer', 'pow', '__pow__', 'prod', 'put_', 'reciprocal', 'remainder', 'repeat', 'repeat_interleave', 'reshape', 'roll', 'round', 'rsqrt', 'scatter', 'scatter_', 'scatter_add', 'select', 'sigmoid', 'sin', 'sinc', 'sinh', 'sort', 'split', 'sqrt', 'square', 'std', 'sub', '__sub__', 'subtract', 'sub_', '__isub__', 'sum', 't', 'take', 'tan', 'tanh', 'tile', 'topk', 'transpose', 'tril', 'triu', 'true_divide', 'trunc', 'type_as', 'unbind', 'unique', 'unsqueeze', 'var', 'view_as', 'where', 'xlogy', '_to']