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

mindspore/_extends/parallel_compile/akg_compiler/gen_custom_op_files.py

@@ -81,7 +81,7 @@ def gen_custom_op_files(config_dir, dsl_dir):
             f.write(json.dumps(ops_info, indent=4))
 
     # custom akg op dsl file
-    custom_py = os.path.join(cur_path, "../../../../../lite/tools/kernel_builder/ascend/akg/custom.py")
+    custom_py = os.path.join(cur_path, "custom.py")
     if not os.path.isfile(custom_py):
         raise RuntimeError("custom.py path is invalid: {}".format(custom_py))
     shutil.copy(custom_py, dsl_dir)

mindspore/_extends/parse/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2024 Huawei Technologies Co., Ltd
+# Copyright 2020-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.
@@ -30,7 +30,7 @@ from .parser import (Parser, create_instance, is_supported_create_instance_type,
                      get_obj_defined_from_obj_type, is_from_third_party_library, get_const_abs, get_const_round,
                      get_const_len, convert_to_namedtuple, check_attrs, generate_lambda_object,
                      check_is_subclass, check_attr_is_property, get_method_info, can_constant_fold,
-                     convert_to_mutable)
+                     convert_to_mutable, get_ast_augassign_namespace_symbol, hook_wrapper)
 
 __all__ = ['Parser', 'create_instance', 'is_supported_create_instance_type', 'generate_scope', 'get_attr_from_object',
            'get_bprop_method_of_class', 'get_class_instance_type', 'get_class_member_namespace_symbol',
@@ -45,4 +45,4 @@ __all__ = ['Parser', 'create_instance', 'is_supported_create_instance_type', 'ge
            'is_class_member_recursive', 'get_obj_defined_from_obj_type',
            'is_from_third_party_library', 'get_const_abs', 'get_const_round', 'get_const_len', 'get_method_info',
            'convert_to_namedtuple', 'check_attrs', 'generate_lambda_object', 'check_is_subclass', 'check_attr_is_property',
-           'can_constant_fold', 'convert_to_mutable']
+           'can_constant_fold', 'convert_to_mutable', 'get_ast_augassign_namespace_symbol', 'hook_wrapper']

mindspore/_extends/parse/compile_config.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Huawei Technologies Co., Ltd
+# Copyright 2024-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.
@@ -275,16 +275,6 @@ Value Range:
 """
 STRICT_CHECK_PARENT_CONTEXT = ''
 
-"""
-Name: CELL_PARAMETER_HOOK
-Function: Whether to enable cell parameter hook.
-          Cell parameter hook is an experimental api that may be deleted later.
-Value Range:
-    1: Enable
-    Default: Disable
-"""
-CELL_PARAMETERS_HOOK = ''
-
 """
 Name: CHECK_BPROP
 Function: Whether to check back propagation nodes. The checking ensures that the shape and dtype of
@@ -315,15 +305,6 @@ Value Range:
 """
 DEBUG_LEVEL = ''
 
-"""
-Name: PYNATIVE_JIT_GRAD_MODE
-Function: Which method used for grad jit in pynative mode
-Value Range:
-    1: Replace ValueNode
-    Default: Parametrization
-"""
-PYNATIVE_JIT_GRAD_MODE = ''
-
 """
 Name: PIJIT_SUBGRAPH_BREAK_OPTIMIZE
 Function: Whether to enable subgraph break optimization in PIJit.
@@ -333,6 +314,15 @@ Value Range:
 """
 PIJIT_SUBGRAPH_BREAK_OPTIMIZE = ''
 
+"""
+Name: ENABLE_ELIMINATE_UNUSED_PARAMS
+Function: Whether to enable eliminate unused parameters optimization in PIJit.
+Value Range:
+    1: Enable, Disable if other value.
+    Default: Disable.
+"""
+ENABLE_ELIMINATE_UNUSED_PARAMS = ''
+
 """
 Name: PUT_ALL_CNODE_INTO_ORDER_LIST
 Function: Whether to put all CNode into order list in back prop.
@@ -342,6 +332,36 @@ Value Range:
 """
 PUT_ALL_CNODE_INTO_ORDER_LIST = ''
 
+"""
+Name: CHECK_PASS_NODE_SCOPE
+Function: Whether to check
+Value Range:
+    1: Enable
+    Default: Disable.
+"""
+CHECK_PASS_NODE_SCOPE = ''
+
+"""
+Name: CHECK_INVALID_VIEW_INPLACE_DOUT_LEVEL
+Function: The level of check invalid dout under view+inplace scene
+Value Range:
+    1: Only check scenario 1
+    2: Only check scenario 2
+    Default(""): Check all invalid dout for view inplace scene
+    Others: No invalid dout check for view inplace scene
+"""
+CHECK_INVALID_VIEW_INPLACE_DOUT_LEVEL = ''
+
+"""
+Name: JIT_ENABLE_AUGASSIGN_INPLACE
+Function: Whether enable augassign inplace.
+Value Range:
+    0: Disable
+    1: Enable
+    Default: Disable
+"""
+JIT_ENABLE_AUGASSIGN_INPLACE = '0'
+
 __all__ = [
     "COMPILE_PROFILE",
     "COMPILE_PROFILE_FINISH_ACTION",
@@ -372,11 +392,13 @@ __all__ = [
     "ENABLE_RECOMPUTE_BEFORE_INLINE",
     "STRICT_CHECK_PARENT_CONTEXT",
     "AUTO_PASSES_OPTIMIZE_PATH",
-    "CELL_PARAMETERS_HOOK",
     "CHECK_BPROP",
     "GRAD_FOR_SCALAR",
     "DEBUG_LEVEL",
-    "PYNATIVE_JIT_GRAD_MODE",
     "PIJIT_SUBGRAPH_BREAK_OPTIMIZE",
+    "ENABLE_ELIMINATE_UNUSED_PARAMS",
     "PUT_ALL_CNODE_INTO_ORDER_LIST",
+    "CHECK_PASS_NODE_SCOPE",
+    "CHECK_INVALID_VIEW_INPLACE_DOUT_LEVEL",
+    "JIT_ENABLE_AUGASSIGN_INPLACE"
 ]

mindspore/_extends/parse/deprecated/deprecated_tensor_method.py

@@ -33,7 +33,7 @@ deprecated_tensor_method_map = {
     # 7 allclose
     "allclose": "tensor_allclose",
     # 8 any
-    "any": "tensor_any",
+    "any": "reduce_tensor_any",
     # 9 arctan2
     "arctan2": "tensor_arctan2",
     # 10 argmax
@@ -338,7 +338,6 @@ deprecated_tensor_method_map = {
     "atan": "deprecated_tensor_atan",
     "arctan": "deprecated_tensor_arctan",
     "dot": "deprecated_tensor_dot",
-    "copy_": "deprecated_tensor_copy_",
 
     # 153
     "logsumexp": "deprecated_tensor_logsumexp",

mindspore/_extends/parse/parser.py

@@ -1,6 +1,6 @@
 # This is the Python adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
 #
-# Copyright 2020-2024 Huawei Technologies Co., Ltd
+# Copyright 2020-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.
@@ -42,9 +42,9 @@ from mindspore.common.api import _JitExecutor
 from mindspore.common import dtype as mstype
 from mindspore.common.parameter import Parameter
 from mindspore.common import mutable
-from mindspore._checkparam import is_stub_tensor
 from .namespace import Namespace, ModuleNamespace, ClosureNamespace, ClassMemberNamespace
-from .resources import parse_object_map, ops_symbol_map, convert_object_map, convert_class_to_function_map, trope_ns
+from .resources import (parse_object_map, parse_augassign_object_map, ops_symbol_map, convert_object_map,
+                        convert_class_to_function_map, trope_ns)
 from .resources import SYMBOL_UNDEFINE, constant_fold_functions
 from .jit_fallback_modules.check_utils import third_party_checker
 from ...common.api import _convert_python_data
@@ -103,17 +103,6 @@ parse_expr_statement_white_list = (
     "append", "insert", "clear", "reverse", "extend", "update", "register_hook",
 )
 
-# Methods that need to reorder after it's caller is used before
-# e.g. We need to reorder `x.register_hook` after x is used in `out = x + 1` when `register_hook` is called.
-# def construct(x):
-#     out = x + 1
-#     x.register_hook(hook_fn)
-#     return out
-# equals to:
-# def construct(x):
-#     x = x.register_hook(hook_fn) # register_hook will return itself when it is called in the graph (in `GRAPH_MODE`).
-#     out = x + 1
-#     return out
 _need_reorder_methods = (
     "register_hook",
 )
@@ -125,10 +114,6 @@ _unsupported_python_builtin_type = (
     set, dict, slice, complex, reversed, type,
 )
 
-# Unsupported python builtin type in JIT Fallback.
-_fallback_unsupported_python_builtin_type = (
-    compile, eval, exec
-)
 
 _global_params = {}
 
@@ -197,10 +182,7 @@ def get_parse_method_of_class(obj, parse_method=None):
     if parse_method is not None:
         method_name = parse_method
     elif isinstance(obj, nn.Cell):
-        if obj._backward_hook:
-            method_name = "_backward_hook_construct"
-        else:
-            method_name = "construct"
+        method_name = "construct"
 
     return get_attr_from_object(obj, method_name)
 
@@ -264,7 +246,7 @@ def resolve_symbol(namespace, symbol):
     if isinstance(resolve_, _JitExecutor):
         logger.debug("Resolve class _JitExecutor, resolve fn instead.")
         resolve_ = resolve_.fn
-    logger.debug(f"Found '{symbol}' in {namespace.__str__()}, resolved: {resolve_} / {type(resolve_)}")
+    logger.debug("Found '%s' in %s, resolved: %s / %s", symbol, namespace, resolve_, type(resolve_))
     return resolve_
 
 
@@ -667,6 +649,14 @@ def get_ast_namespace_symbol(obj):
     return ops_info
 
 
+def get_ast_augassign_namespace_symbol(obj):
+    """Get obj type and namespace and symbol."""
+    # Get symbol from object map.
+    ops_info = parse_augassign_object_map.get(type(obj), SYMBOL_UNDEFINE)
+    logger.debug("ops info: %r", ops_info)
+    return ops_info
+
+
 def get_operation_symbol(obj):
     """Get obj operation symbol."""
     ops_symbol = ops_symbol_map.get(type(obj), SYMBOL_UNDEFINE)
@@ -680,6 +670,7 @@ def get_operation_namespace_symbol(var: str):
     logger.debug("get operation ops info: %r", ops_info)
     return ops_info
 
+
 def get_ast_type(node):
     """Get the ast type."""
     ast_type = AST_SUB_TYPE_UNKNOWN
@@ -784,32 +775,6 @@ def get_arg_spec_and_default_values(func):
     return arg_spec, defaults
 
 
-def _convert_stub_tensor(data):
-    """Convert stub tensor output to tensor"""
-    if is_stub_tensor(data):
-        return data.stub_sync()
-    if isinstance(data, tuple):
-        # Handle namedtuple since its type is tuple.
-        if hasattr(data, "_fields"):
-            type_name = data.__class__.__name__
-            data_dict = data._asdict()
-            fields = data_dict.keys()
-            return namedtuple(type_name, fields)(**_convert_stub_tensor(data_dict))
-        return tuple(_convert_stub_tensor(x) for x in data)
-    if data.__class__ is list:
-        # Keep the list object not change.
-        for i in range(len(data)):
-            data[i] = _convert_stub_tensor(data[i])
-        return data
-    if data.__class__ is dict:
-        # Keep the dict object not change.
-        keys = tuple(data.keys())
-        for key in keys:
-            data[_convert_stub_tensor(key)] = _convert_stub_tensor(data.pop(key))
-        return data
-    return data
-
-
 def eval_script(exp_str, params):
     """Evaluate a python expression."""
     if not isinstance(params, tuple):
@@ -822,8 +787,13 @@ def eval_script(exp_str, params):
     local_params = params[1]
     try:
         local_params = _convert_python_data(local_params)
+        # There are two sources of scripts:
+        # 1. The user's original Python script code, which is directly passed back to Python for execution,
+        #    and its behavior is guaranteed by the user.
+        # 2. Internally provided Python expression code, similar to
+        #    `__iternal_sequence_input__[__internal_sequence_index__]`.
+        # In addition, MindIR load and export do not involve the use of the `eval_script` function.
         res = eval(exp_str, global_params, local_params)
-        res = _convert_stub_tensor(res)
     except Exception as e:
         error_info = f"When eval '{exp_str}' by using JIT Fallback feature, an error occurred: " + str(e)
         logger.debug(error_info)
@@ -852,7 +822,7 @@ def get_script_id_attrs(script):
 
 def generate_lambda_object(script):
     """Generate lambda expression object using script"""
-    return eval(script, {}, {})
+    return eval_script(script, ({}, {}))
 
 
 def get_global_params():
@@ -924,25 +894,31 @@ def get_method_info(obj):
     return class_name_and_method_name[0], class_name_and_method_name[1]
 
 
-def is_ms_tensor_method(obj):
-    """Check if the obj is a method of MindSpore Tensor"""
-    if not hasattr(obj, '__name__') or not hasattr(Tensor, obj.__name__):
-        return False
-    fn = inspect.unwrap(obj.__func__ if isinstance(obj, types.MethodType) else obj)
-    tensor_method = getattr(Tensor, obj.__name__)
-    tensor_method = tensor_method.__func__ if hasattr(tensor_method, "__func__") else tensor_method
-    is_builtin_tensor = False
-    if inspect.isbuiltin(obj):
-        class_name_and_method_name = obj.__qualname__.split('.')
-        is_builtin_tensor = obj.__module__ is None and class_name_and_method_name[0] == Tensor.__name__
-    return fn == tensor_method or is_builtin_tensor
-
-
 def can_constant_fold(obj):
     """Check if the obj is the function can be constantly folded."""
     return obj in constant_fold_functions
 
 
+def hook_wrapper(hook_fn):
+    """
+    Decorator wrapper for gradient hook functions.
+    Handles custom logic when the hook returns None to ensure execution dependencies.
+
+    Args:
+        hook_fn (function): The original hook function to be wrapped.
+
+    Returns:
+        function: Wrapped inner hook function with dependency handling logic.
+    """
+    def inner(dout):
+        fdout = hook_fn(dout)
+        if fdout is None:
+            dout = ops.Depend()(dout, fdout)
+            return dout
+        return fdout
+    return inner
+
+
 class Parser:
     """
     Parser python code to ast tree.
@@ -976,8 +952,6 @@ class Parser:
         """To check if not supported for namespace"""
         unsupported = isinstance(value, _builtin_function_or_method_type) and value not in convert_object_map
         logger.debug(f"'{value}' unsupported: {unsupported}.")
-        if unsupported and value in _fallback_unsupported_python_builtin_type:
-            raise TypeError(f"'{value}' is not supported both in JIT Fallback and graph mode.")
         return unsupported
 
     @staticmethod
@@ -1037,31 +1011,38 @@ class Parser:
                              "the code 'def __init__(self, combine_fn=lambda x: x + 1):' rewritten as\n"
                              "'def __init__(self, combine_fn=\nlambda x: x + 1\n):' will solve the problem.")
 
+    def save_source_code(self, attr_name, source_lines):
+        """Save cell and func source code to support run graph mode with pyc or so."""
+        if '/mindspore/' in self.filename or '\\mindspore\\' in self.filename:
+            return
+        if getattr(self.fn, attr_name, None) == source_lines:
+            return
+        if not os.access(self.filename, os.W_OK):
+            raise PermissionError(f"Don't have the write permission on the file {self.filename}.")
+        with open(self.filename, 'a') as f:
+            logger.debug(f"setattr for {self.fn}, attr: {attr_name}, value: {source_lines}")
+            f.write(f"\n# Set source attribute for function {self.function_name} "
+                    f"to support run so or pyc file in Graph Mode."
+                    f"\nsetattr({self.function_name}, '{attr_name}', {source_lines})\n")
+            setattr(self.fn, attr_name, source_lines)
+
     def parse(self):
         """Parse the function or method."""
         logger.debug("fn: %r", self.fn)
         if isinstance(self.fn, (types.FunctionType, types.MethodType)) or \
            type(self.fn).__name__ == 'cython_function_or_method':
-            attr = 'source'
+            attr_name = 'source'
             try:
                 source_lines = inspect.getsourcelines(self.fn)
-                if context.get_context('support_binary') and \
-                   '/mindspore/' not in self.filename and '\\mindspore\\' not in self.filename and \
-                   (not hasattr(self.fn, attr) or getattr(self.fn, attr) != source_lines):
-                    if not os.access(self.filename, os.W_OK):
-                        raise PermissionError(f"Don't have the write permission on the file {self.filename}.")
-                    with open(self.filename, 'a') as f:
-                        f.write(f"\n# Set source attribute for function {self.function_name} "
-                                f"to support run so or pyc file in Graph Mode."
-                                f"\nsetattr({self.function_name}, '{attr}', {source_lines})\n")
-                        setattr(self.fn, attr, source_lines)
+                if context.get_context('support_binary') or os.getenv('MS_SUPPORT_BINARY', None) == '1':
+                    self.save_source_code(attr_name, source_lines)
             except (OSError, TypeError) as e:
-                if hasattr(self.fn, attr):
-                    source_lines = getattr(self.fn, attr)
+                if hasattr(self.fn, attr_name):
+                    source_lines = getattr(self.fn, attr_name)
+                elif e.__str__() == "could not get source code":
+                    raise OSError(f"Mindspore can not compile temporary source code in terminal. "
+                                  f"Please write source code to a python file and run the file.")
                 else:
-                    if e.__str__() == "could not get source code":
-                        raise OSError(f"Mindspore can not compile temporary source code in terminal. "
-                                      f"Please write source code to a python file and run the file.")
                     raise e
             self.lines, self.line_offset = source_lines
             original_src = ''.join(self.lines)
@@ -1072,7 +1053,7 @@ class Parser:
                 self.col_offset = \
                     len(original_src.split('\n')[0]) - len(src.split('\n')[0])
                 logger.debug("Get source: %s", src)
-                if not hasattr(self.fn, attr):
+                if not hasattr(self.fn, attr_name):
                     self.check_lambda(src)
                 try:
                     ast_tokens = asttokens.ASTTokens(src, parse=True)

mindspore/_extends/parse/resources.py

@@ -76,6 +76,45 @@ parse_object_map = {
     SYMBOL_UNDEFINE: (None, 'undefine', ''),
 }
 
+parse_augassign_object_map = {
+    # ast grammar
+    ast.Add:        (trope_ns, 'augassign_add', '+='),
+    ast.Sub:        (trope_ns, 'augassign_sub', '-='),
+    ast.Mult:       (trope_ns, 'augassign_mul', '*='),
+    ast.Div:        (trope_ns, 'augassign_div', '/='),
+    ast.FloorDiv:   (trope_ns, 'augassign_floordiv', '//='),
+
+    ast.Mod:        (trope_ns, 'mod', '%'),
+    ast.Pow:        (trope_ns, 'pow', '**'),
+    ast.MatMult:    (trope_ns, 'matmul', '@'),
+    ast.LShift:     (trope_ns, 'lshift', '<<'),
+    ast.RShift:     (trope_ns, 'rshift', '>>'),
+    ast.BitAnd:     (trope_ns, 'and_', '&'),
+    ast.BitOr:      (trope_ns, 'or_', '|'),
+    ast.BitXor:     (trope_ns, 'xor', '^'),
+    ast.UAdd:       (trope_ns, 'pos', '+'),
+    ast.USub:       (trope_ns, 'neg', '-'),
+    ast.Invert:     (trope_ns, 'invert', '~'),
+    ast.Not:        (trope_ns, 'not_', 'not'),
+    ast.Eq:         (trope_ns, 'eq', '=='),
+    ast.NotEq:      (trope_ns, 'ne', '!='),
+    ast.Lt:         (trope_ns, 'lt', '<'),
+    ast.Gt:         (trope_ns, 'gt', '>'),
+    ast.LtE:        (trope_ns, 'le', '<='),
+    ast.GtE:        (trope_ns, 'ge', '>='),
+    ast.Is:         (trope_ns, 'is_', 'is'),
+    ast.IsNot:      (trope_ns, 'is_not', 'is not'),
+    ast.In:         (trope_ns, 'contains', 'in'),
+    ast.NotIn:      (trope_ns, 'not_contains', 'not in'),
+
+    # operation symbol type
+    'getitem':      (composite_ns, 'getitem', ''),
+    'ms_next':      (composite_ns, 'ms_next', ''),
+
+    # undefined type
+    SYMBOL_UNDEFINE: (None, 'undefine', ''),
+}
+
 # Operation symbols corresponding to ast grammar
 ops_symbol_map = {
     # ast grammar

mindspore/_extends/parse/standard_method.py

@@ -27,9 +27,11 @@ from mindspore.common.sparse_tensor import RowTensorInner
 from mindspore.ops.composite.base import _append, _insert, _pop, _list_clear, _reverse, \
     _extend, _dict_setitem, _dict_clear, _haskey, _update, _fromkeys
 from mindspore.ops.operations._sequence_ops import TensorToTuple
-from mindspore.ops.auto_generate import trace_v2_op, inplace_addmm_op, inplace_index_put_op, inplace_normal_op, inplace_index_add_op
+from mindspore.ops.auto_generate import trace_v2_op, inplace_addmm_op, inplace_index_put_op, inplace_normal_op, \
+    inplace_index_add_op
 from mindspore.ops.auto_generate import inplace_copy_op, inplace_uniform_op, inplace_erfinv_op
 from mindspore.ops.auto_generate import inplace_scatter_add as inplace_scatter_add_
+from mindspore.ops.auto_generate import inplace_exponential_op
 
 from ... import _checkparam as validator
 from ..._checkparam import check_is_number, check_reshape_shp, check_axis_in_range, \
@@ -579,7 +581,7 @@ def transpose(x, *axis):
 
     Raises:
         TypeError: If input arguments have types not specified above.
-        ValueError: If the number of `axes` is not euqal to a.ndim.
+        ValueError: If the number of `axes` is not equal to a.ndim.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -1112,7 +1114,7 @@ def copy_(self, src, non_blocking=False):
     """
     Copies the elements from src into self tensor and returns self.
     """
-    return inplace_copy_op(self, src)
+    return inplace_copy_op(self, src, non_blocking)
 
 
 def max(input, axis=None, keepdims=False, *, initial=None,  # pylint: disable=redefined-builtin
@@ -1232,7 +1234,7 @@ def pow(x, y):  # pylint: disable=redefined-builtin
     return F.pow(x, y)
 
 
-def put_(x, index, source, accumulate=False):  # pylint: disable=redefined-builtin
+def put_(x, index, source, accumulate=False):
     """
     Copies the elements from source into the positions specified by index.
     """
@@ -2131,14 +2133,14 @@ def _check_sum_to_size(size, input_dim, shape_input):
 
 
 @_primexpr
-def _count_axes(size, input_shape, shape_input):
+def _count_axes(size, input_shape, shape_input, pre_len, pre_axis):
     """Count the sum axes for sum_to_size."""
-    axes = []
+    axes = pre_axis
     for i in range(len(size)):
         element = size[i]
-        if element != input_shape[i] and element == 1:
-            axes.append(i)
-        elif element != input_shape[i]:
+        if element != input_shape[i + pre_len] and element == 1:
+            axes.append(i + pre_len)
+        elif element != input_shape[i + pre_len]:
             raise ValueError(f"For sum_to_size, size {size} is not expandable to the tensor size {shape_input}.")
     return axes
 
@@ -2151,13 +2153,15 @@ def sum_to_size(input, *size):
         size = size[0]
     shape_input = input.shape
     _check_sum_to_size(size, input.ndim, shape_input)
+    pre_len = 0
+    pre_axis = []
     if len(size) < input.ndim:
-        pre_axis = tuple(axis for axis in range(input.ndim - len(size)))
-        input = input.sum(pre_axis)
+        pre_len = input.ndim - len(size)
+        pre_axis = [axis for axis in range(pre_len)]
 
-    axes = _count_axes(size, input.shape, shape_input)
+    axes = _count_axes(size, input.shape, shape_input, pre_len, pre_axis)
     if axes:
-        return input.sum(tuple(axes), keepdims=True)
+        return input.sum(tuple(axes), keepdims=True).reshape(size)
     return input
 
 
@@ -3705,6 +3709,13 @@ def bernoulli(input, *, generator=None):
     return F.bernoulli_ext(input, generator=generator)
 
 
+def bernoulli_(input, p=0.5, *, generator=None):
+    """
+    Randomly draws binary numbers from a Bernoulli distribution.
+    """
+    return F.bernoulli_(input, p, generator=generator)
+
+
 def gather_nd(input_x, indices):
     r"""
     Gathers slices from a tensor by indices.
@@ -4004,6 +4015,7 @@ def to_double(input_x):
     """
     return F.cast(input_x, mstype.float64)
 
+
 def to_bfloat16(input_x):
     r"""
     Converts input tensor dtype to bfloat16.
@@ -4486,10 +4498,20 @@ def uniform_(input, from_=0, to=1, *, generator=None):
     """
     if generator is None:
         generator = default_generator
-    seed, offset = generator._step(generator_step_) # pylint: disable=protected-access
+    seed, offset = generator._step(generator_step_)  # pylint: disable=protected-access
     return inplace_uniform_op(input, from_, to, seed, offset)
 
 
+def exponential_(input, lambd=1, *, generator=None):
+    r"""
+        Fills `self` tensor with elements drawn from the exponential distribution:
+    """
+    if generator is None:
+        generator = default_generator
+    seed, offset = generator._step(generator_step_)  # pylint: disable=protected-access
+    return inplace_exponential_op(input, lambd, seed, offset)
+
+
 def amin(input, axis=None, keep_dims=False):
     r"""
     For details, please refer to :func:`mindspore.ops.amin`.
@@ -4520,6 +4542,13 @@ def index_put(input, indices, values, accumulate=False):
     return _index_put(input, values, indices)
 
 
+def move_to(input, to, blocking=True):
+    r"""
+        Copy Tensor to target device synchronously or asynchronously, default synchronously. only support PyNative mode.
+    """
+    raise ValueError(f"The method 'move_to' is not supported in jit.")
+
+
 def index_put_(input, indices, values, accumulate=False):
     r"""
     For details, please refer to :func:`mindspore.Tensor.index_put_`.
@@ -4592,3 +4621,18 @@ def zero_(input):
     Return a tensor filled with zeros.
     """
     return F.zero_(input)
+
+
+def slice_get_start(slice_node):
+    """Using SliceGetItem to get slice_node.start"""
+    return F.SliceGetItem(slice_node, "start")
+
+
+def slice_get_stop(slice_node):
+    """Using SliceGetItem to get slice_node.stop"""
+    return F.SliceGetItem(slice_node, "stop")
+
+
+def slice_get_step(slice_node):
+    """Using SliceGetItem to get slice_node.step"""
+    return F.SliceGetItem(slice_node, "step")

mindspore/_extends/parse/trope.py

@@ -37,6 +37,12 @@ from functools import (  # noqa
     partial
 )
 
+from mindspore.ops.composite.multitype_ops.add_impl import augassign_add
+from mindspore.ops.composite.multitype_ops.sub_impl import augassign_sub
+from mindspore.ops.composite.multitype_ops.mul_impl import augassign_mul
+from mindspore.ops.composite.multitype_ops.div_impl import augassign_div
+from mindspore.ops.composite.multitype_ops.floordiv_impl import augassign_floordiv
+
 from ...common import mutable
 
 __all__ = ['add', 'sub', 'mul', 'truediv', 'floordiv', 'mod', 'eq', 'ne', 'lt', 'gt', 'le', 'ge', 'pos', 'neg',
@@ -44,7 +50,8 @@ __all__ = ['add', 'sub', 'mul', 'truediv', 'floordiv', 'mod', 'eq', 'ne', 'lt',
            'matmul', 'getitem', 'setitem',
            'bool', 'getattr', 'setattr', 'hasattr', 'len', 'iter', 'next', 'pow', 'range', 'map', 'zip',
            'partial', 'print', 'enumerate', 'isinstance', 'filter', 'abs', 'round', 'mutable',
-           'max', 'min', 'sum', 'list', 'tuple']
+           'max', 'min', 'sum', 'list', 'tuple',
+           'augassign_add', 'augassign_sub', 'augassign_mul', 'augassign_div', 'augassign_floordiv']
 
 
 def MakeTuple(*elts):  # pragma: no cover

mindspore/_extends/pijit/__init__.py

@@ -16,8 +16,7 @@
 Helper module for pijit analyze
 """
 
+__all__ = ['pijit_func_white_list_map', "get_tensor_method_name"]
 
 from .pijit_func_white_list import _func_map as pijit_func_white_list_map
 from .tensor_func_list import get_tensor_method_name
-
-__all__ = ['pijit_func_white_list_map', "get_tensor_method_name"]