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

mindspore/common/_grad_function.py

@@ -0,0 +1,56 @@
+# Copyright 2025 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.
+# ============================================================================
+
+"""Defines custom autograd function with functional form."""
+
+from typing import Any
+from mindspore._c_expression import FunctionBase as FunctionBase_
+from mindspore.common.tensor import Tensor
+
+__all__ = ['_Function']
+
+class _Function(FunctionBase_):
+    """
+    A Class provides the ability to custom autograd function.
+
+    Note:
+        It is only supported in pynative mode.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+    """
+    @staticmethod
+    def forward(ctx: Any, *args: Any, **kwars: Any) -> Any:
+        raise NotImplementedError("forward function should be customized.")
+
+    @staticmethod
+    def backward(ctx: Any, *grad_outputs: Any) -> Any:
+        raise NotImplementedError("backward function should be customized.")
+
+    @classmethod
+    def apply(cls, *args, **kwargs):
+        return super().apply(cls, *args, **kwargs)
+
+    def save_for_backward(self, *tensors: Tensor):
+        self.saved_tensors = tensors
+
+    def mark_dirty(self, *args: Tensor):
+        self.dirty_tensors = args
+
+    def mark_non_differentiable(self, *args: Tensor):
+        self.non_differentiable = args
+
+    def set_materialize_grads(self, value: bool):
+        self.materialize_grads = value

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:

mindspore/common/_stub_tensor.py

@@ -18,8 +18,9 @@ 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
 
@@ -57,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
@@ -177,7 +172,7 @@ class StubTensor:
         """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)
@@ -193,7 +188,7 @@ 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"]
 
@@ -211,7 +206,7 @@ def _init_stub_tensor_api():
 
 
 _init_stub_tensor_api()
-
+c_expression.addStubTensorMethods()
 
 def _convert_stub(stub):
     "convert stub to StubNode or Value"

mindspore/common/_tensor_cpp_method.py

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