mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

mindspore/ops/_tracefunc.py

@@ -0,0 +1,241 @@
+# 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.
+# ============================================================================
+"""Context for PackFunc"""
+import functools
+import types
+import textwrap
+import inspect
+import os
+from mindspore.common.tensor import Tensor
+from mindspore.ops.primitive import _RunOpHook, Primitive
+from mindspore._c_expression import PackExpander, PackNode
+from mindspore.common._stub_tensor import StubTensor
+from mindspore.common._register_for_tensor import tensor_operator_registry
+from mindspore.common.api import _handle_func_args, _pynative_executor
+
+
+class _PackTensor(StubTensor):
+    """pack tensor for expander trace"""
+
+    def __setitem__(self, index, value):
+        out = tensor_operator_registry.get('__setitem__')(self, index, value)
+        self.stub = out.stub
+        if self.parent_tensor_ is not None and self.index_of_parent_ is not None:
+            self.parent_tensor_.__setitem__(self.index_of_parent_, self)
+        return self
+
+    def __pack__(self):
+        """for parse check."""
+
+    def stub_sync(self):
+        """subclass hook for Tensor"""
+        if self.tensor is None:
+            val = self.stub.get_value()
+            if val is None:
+                raise RuntimeError("During the trace operation, the data flow of the Tensor could be tracked, "
+                                   "which consequently prevented the creation of a proper trace subgraph.")
+            self.tensor = Tensor(val, internal=True)
+            if hasattr(self, "member_cache"):
+                for k, v in self.member_cache.items():
+                    setattr(self.tensor, k, v)
+        return self.tensor
+
+
+def _convert_tensor(node):
+    if isinstance(node, PackNode):
+        return _PackTensor(node)
+    if isinstance(node, tuple):
+        return tuple(_convert_tensor(e) for e in node)
+    return node
+
+
+class PackFunc(Primitive):
+    """pack function with lazy expander"""
+
+    expander = PackExpander.get_instance()
+    current = None
+
+    def __init__(self, fun, unique_key, cell_obj, is_pynative_mode=False):
+        super(PackFunc, self).__init__(self.__class__.__name__)
+        self.func = fun
+        self.cell_obj = cell_obj
+        self.kwargs = {}
+        self.add_prim_attr("unique_key", unique_key)
+        self.add_prim_attr("is_pynative_mode", is_pynative_mode)
+
+    def __call__(self, *args, **kwargs):
+        if PackFunc.is_tracing():
+            if self.cell_obj:
+                args = (self.cell_obj, *args)
+            return self.func(*args, **kwargs)
+        self.kwargs = kwargs
+        output = super().__call__(*args)
+        if self.is_pynative_mode and self.grad_attach_num > 0:
+            output_num = len(output) - self.grad_attach_num
+            if output_num == 1:
+                return output[0]
+            return output[:output_num]
+        return output
+
+    def __expand__(self, args):
+        old = PackFunc.current
+        PackFunc.current = self
+        if self.cell_obj:
+            args = (self.cell_obj, *args)
+            with _SetMixedPrecision(self.cell_obj):
+                ret = self._run_op(args)
+        else:
+            ret = self._run_op(args)
+        PackFunc.current = old
+        return ret
+
+    @staticmethod
+    def is_tracing():
+        return PackFunc.current is not None
+
+    @staticmethod
+    def _trace_run_op(obj, args):
+        ret = PackFunc.expander.emit(obj, *args)
+        return _convert_tensor(ret)
+
+    def _run_op(self, args):
+        with _RunOpHook(PackFunc._trace_run_op):
+            fun_args = [_convert_tensor(a) for a in args]
+            ret = self.func(*fun_args, **self.kwargs)
+        return ret
+
+
+class _PackSourceBuilder:
+    """Generation Pack Python code by method"""
+
+    def __init__(self, original_fn):
+        """Initialize the _PackSourceBuilder"""
+        self.original_fn = original_fn
+        self.is_method = isinstance(self.original_fn, types.MethodType)
+        self.pack_fn_name = f"_{original_fn.__name__}_pack"
+        self._generate_pack_op()
+
+    def get_code_source(self):
+        """Return Pack Python code"""
+        if self.is_method:
+            sig = inspect.signature(self.original_fn.pack_fn)
+            arg_num = len(sig.parameters) - 1
+            arg_str = ", ".join(["a{}".format(i) for i in range(arg_num)])
+            new_src = textwrap.dedent(f"""
+                def {self.pack_fn_name}_wrap(self, {arg_str}):
+                    return self.{self.pack_fn_name}({arg_str})
+                """)
+        else:
+            new_src = textwrap.dedent(f"""
+                def {self.pack_fn_name}_wrap(*args, **kwargs):
+                    return {self.pack_fn_name}(*args, **kwargs)
+                """)
+        return new_src
+
+    def _generate_pack_op(self):
+        """Generate the pack operation and attach it to the original function or method"""
+        obj = self.original_fn.__self__ if self.is_method else None
+        fn = self.original_fn.pack_fn
+        key = f"{id(obj)}_{id(fn)}"
+        if self.is_method:
+            setattr(obj, self.pack_fn_name, PackFunc(fn, key, obj))
+        else:
+            fn.__globals__[self.pack_fn_name] = PackFunc(fn, key, obj)
+
+
+class _SetMixedPrecision:
+    """"Set MixedPrecison by the Cell"""
+
+    def __init__(self, cell_obj):
+        self.mixed_precision_change = False
+        self.cell_obj = cell_obj
+
+    def __enter__(self):
+        self.mixed_precision_change = PackFunc.expander.set_mixed_precision(
+            self.cell_obj)
+
+    def __exit__(self, *err):
+        if self.mixed_precision_change:
+            PackFunc.expander.recover_mixed_precision()
+
+
+def trace(fn):
+    """
+    Create a traceable function from a python function. The python function will be traced with fake tensor to
+    capture the corresponding runtime graph at compilation phase.Tracing is ideal for some scenarios to reduce
+    python overhead or graph compilation time.
+
+    Note:
+        - Dynamic control flow and other tensor data dependent scenarios are not supported.
+        - Dynamic shape operator is not fully supported.
+
+    Args:
+        fn (Function) - The Python function to be compiled into a graph. The return value of the function must be
+        a Tensor or a Tuple containing Tensors. When a network's construct is traced, only its forward method is traced.
+
+    Returns:
+        Function, if `fn` is not None, returns a callable function that will execute the compiled function; If `fn` is
+        None, returns a decorator and when this decorator invokes with a single `fn` argument, the callable function is
+        equal to the case when `fn` is not None.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> from mindspore import Tensor, nn
+        >>> from mindspore.ops._tracefunc import trace
+        ...
+        >>> class Net(nn.Cell):
+        ...     @trace
+        ...     def construct(self, x, y):
+        ...         return x + y
+        >>> x = Tensor([1, 1, 3, 3])
+        >>> y = Tensor([1, 1, 3, 3])
+        >>> net = Net()
+        >>> output = net(x, y)
+    """
+
+    @functools.wraps(fn)
+    def _trace_wrap(*args, **kwargs):
+        pynative_grad_flag = _pynative_executor.grad_flag()
+        grad_flag_expr = "1" if pynative_grad_flag else "0"
+        if _trace_wrap.is_method is None:
+            if args and not isinstance(args[0], Tensor) and hasattr(args[0], fn.__name__):
+                _trace_wrap.is_method = False
+            else:
+                _trace_wrap.is_method = True
+        if _trace_wrap.is_method:
+            # Similar processing has been done in the __call__ of Cell,
+            # so only when obj is None, there is need to do `_handle_func_args`.
+            args, kwargs = _handle_func_args(fn, *args, **kwargs)
+            pack_func_name = "pack" + grad_flag_expr
+            pack_func = getattr(fn, pack_func_name, None)
+            if pack_func is None:
+                pack_func = PackFunc(fn, f"{id(fn)}_{grad_flag_expr}", None, True)
+                setattr(fn, pack_func_name, pack_func)
+            return pack_func(*args, **kwargs)
+        obj, args = args[0], args[1:]
+        pack_func_name = "".join((fn.__name__, "pack", grad_flag_expr))
+        pack_func = getattr(obj, pack_func_name, None)
+        if pack_func is None:
+            pack_func = PackFunc(fn, f"{id(obj)}_{id(fn)}_{grad_flag_expr}", obj, True)
+            setattr(obj, pack_func_name, pack_func)
+        return pack_func(*args, **kwargs)
+
+    if "MS_DEV_DISABLE_TRACE" in os.environ and os.environ["MS_DEV_DISABLE_TRACE"] == "on":
+        return fn
+    _trace_wrap.pack_fn = fn
+    _trace_wrap.is_method = None
+    return _trace_wrap

mindspore/ops/_utils/utils.py

@@ -77,12 +77,15 @@ def get_broadcast_shape(x_shape, y_shape, prim_name, arg_name1="x", arg_name2="y
     broadcast_shape = list(broadcast_shape_front) + broadcast_shape_back
     return broadcast_shape
 
+def dim_not_equal(dim1, dim2):
+    """Compare dim in shape"""
+    return dim1 != dim2 and dim1 >= 0 and dim2 >= 0
 
 def get_concat_offset(x_shp, x_type, axis, prim_name):
     """for concat and concatoffset check args and compute offset"""
     validator.check_value_type("shape", x_shp, [tuple, list], prim_name)
     validator.check_positive_int(len(x_shp), "input_x rank", prim_name)
-    validator.check_subclass("shape0", x_type[0], mstype.tensor, prim_name)
+    validator.check_subclass("shape0", x_type[0], mstype.tensor_type, prim_name)
     validator.check_positive_int(len(x_shp[0]), "len of x_shp[0]", prim_name)
     rank_base = len(x_shp[0])
     for i in range(1, len(x_shp)):
@@ -98,7 +101,7 @@ def get_concat_offset(x_shp, x_type, axis, prim_name):
     for i in range(1, len(x_shp)):
         v = x_shp[i]
         for j in range(rank_base):
-            if j != axis and v[j] != x_shp[0][j] and v[j] >= 0 and x_shp[0][j] >= 0:
+            if j != axis and dim_not_equal(v[j], x_shp[0][j]):
                 raise ValueError(f"The shape of the two input elements of the Concat operator do not match:"
                                  f"shape[0] = {x_shp[0]} and shape[{i}] = {x_shp[i]}.")
         offset.append(all_shp)
@@ -135,3 +138,8 @@ def generate_shape_index(out_shape, indices_shape, axis, batch_dims=0):
     index = tuple(range(out_rank))
     perm = index[:batch_dims] + perm_part1 + index[batch_dims:axis] + index[axis + ind_rank - batch_dims:]
     return perm
+
+
+def ms_arrange(x):
+    out = [i for i in range(x)]
+    return out

mindspore/ops/_vmap/vmap_array_ops.py

@@ -155,7 +155,7 @@ def _get_prefix(indices_shape, axis_size, indices_dtype):
     _check(indices_shape)
     indices_len = len(indices_shape)
     if indices_len == 1:
-        prefix = P.Range()(Tensor(0, indices_dtype), P.Fill()(
+        prefix = P.Range()(Tensor(0, indices_dtype), F.fill(
             indices_dtype, (), axis_size), Tensor(1, indices_dtype))
         return prefix
 
@@ -850,9 +850,9 @@ def get_fill_vmap_rule(prim, axis_size):
 
 
 @constexpr
-def to_tensor_with_type(x, type):
+def to_tensor_with_type(x, dtype):
     """x to Tensor with type"""
-    return Tensor(x, type)
+    return Tensor(x, dtype)
 
 
 @vmap_rules_getters.register(P.FillV2)
@@ -1138,6 +1138,8 @@ def get_masked_select_vmap_rule(prim, axis_size):
 @vmap_rules_getters.register(MaskedSelectGrad)
 def get_masked_select_grad_vmap_rule(prim, axis_size):
     """VmapRule for `MaskedSelect`."""
+    if isinstance(prim, str):
+        prim = Primitive(prim)
 
     def vmap_rule(x_bdim, mask_bdim, outgrad_bdim):
         is_all_none, result = vmap_general_preprocess(prim, x_bdim, mask_bdim, outgrad_bdim)

mindspore/ops/_vmap/vmap_base.py

@@ -28,7 +28,7 @@ from mindspore.ops.operations import nn_ops as nps
 from mindspore.ops.function import _VmapGeneralPreprocess
 from mindspore.ops.primitive import Primitive, _PrimitiveC
 from mindspore.ops.operations.random_ops import UniformCandidateSampler, RandomShuffle
-from mindspore.ops._grad.grad_base import BpropRegistry as VmapRuleRegistry
+from mindspore.ops._grad_experimental.grad_base import BpropRegistry as VmapRuleRegistry
 
 
 vmap_rules_getters = VmapRuleRegistry()
@@ -250,7 +250,7 @@ def vmap_monad_rule(prim, axis_size):
 def _bdim_at_any(x, src, dst, axis_size):
     """
     Moves source axes of an array to the dst axis, and other axes remain in their original order. If the source axes
-    is 'None', broadcasts the array at dst axis with axis_size.
+    is ``None``, broadcasts the array at dst axis with axis_size.
 
     Args:
         x (Tensor or Scalar): The input tensor or scalar. The data type should be one of the following types: float16,
@@ -272,7 +272,7 @@ def _bdim_at_any(x, src, dst, axis_size):
 def _bdim_at_front(x, src, axis_size):
     """
     Moves source axes of an array to the foremost, and other axes remain in their original order. If the source axes
-    is 'None', broadcasts the array at foremost axis with axis_size.
+    is ``None``, broadcasts the array at foremost axis with axis_size.
 
     Args:
         x (Tensor or Scalar): The input tensor or scalar. The data type should be one of the following types: float16,
@@ -289,7 +289,7 @@ def _bdim_at_front(x, src, axis_size):
 def _bdim_at_back(x, src, axis_size):
     """
     Moves source axes of an array to the last, and other axes remain in their original order. If the source axes
-    is 'None', broadcasts the array at foremost axis with axis_size.
+    is ``None``, broadcasts the array at foremost axis with axis_size.
 
     Args:
         x (Tensor or Scalar): The input tensor or scalar. The data type should be one of the following types: float16,
@@ -499,6 +499,7 @@ _ops_vmap_clone_prim_dict = {
     "ApplyProximalAdagrad": P.ApplyProximalAdagrad,
     "ApplyProximalGradientDescent": P.ApplyProximalGradientDescent,
     "ApplyAdamWithAmsgrad": P.ApplyAdamWithAmsgrad,
+    "ApplyAdamWithAmsgradV2": P.ApplyAdamWithAmsgradV2,
     "ApplyPowerSign": P.ApplyPowerSign,
     "ApplyAdagradDA": P.ApplyAdagradDA,
     "ApplyAdagradV2": P.ApplyAdagradV2,

mindspore/ops/_vmap/vmap_convolution_ops.py

@@ -190,8 +190,8 @@ def _reshape_expand_dims(src_dim, dst_size, target, prim_name):
 @_primexpr
 def _get_new_size_by_index(input_size, batch_size, index):
     """Get the new size of input_size by multiplying input_size[index] by batch_size."""
-    new_size = ()
     if input_size is None:
+        new_size = ()
         return new_size
     new_size = list(input_size)
     new_size[index] *= batch_size

mindspore/ops/_vmap/vmap_grad_math_ops.py

@@ -62,8 +62,9 @@ def get_broadcast_binary_op_grad_vmap_rule(prim, axis_size):
         y_shape = F.shape(y)
         g_shape = F.shape(g)
 
-        if x_dim == y_dim and x_dim == g_dim and \
-            x_shape == y_shape and x_shape == g_shape:
+        is_dim_ok = x_dim == y_dim and x_dim == g_dim
+        is_shape_ok = x_shape == y_shape and x_shape == g_shape
+        if is_dim_ok and is_shape_ok:
             dx, dy = prim(x, y, g)
             return (dx, x_dim), (dy, y_dim)
 
@@ -113,8 +114,9 @@ def get_broadcast_grad_grad_vmap_rule(prim, axis_size):
         dx1_shape = F.shape(dx1)
         dx2_shape = F.shape(dx2)
 
-        if x1_dim == x2_dim and dx1_dim == dx2_dim and x1_dim == dx1_dim \
-                and x1_shape == x2_shape and dx1_shape == dx2_shape:
+        is_dim_ok = x1_dim == x2_dim and dx1_dim == dx2_dim and x1_dim == dx1_dim
+        is_shape_ok = x1_shape == x2_shape and dx1_shape == dx2_shape
+        if is_dim_ok and is_shape_ok:
             sopd_x1, sopd_x2, sopd_grad = prim(x1, x2, dx1, dx2)
             return (sopd_x1, x1_dim), (sopd_x2, x1_dim), (sopd_grad, x1_dim)
 

mindspore/ops/_vmap/vmap_grad_nn_ops.py

@@ -26,7 +26,7 @@ from mindspore.ops.primitive import _primexpr
 from mindspore.ops.primitive import Primitive
 from mindspore.ops.function import _VmapGeneralRule
 from mindspore.ops._vmap.vmap_base import vmap_rules_getters, vmap_general_preprocess, _raise_value_error, \
-    _bdim_at_front, _vmap_clone_prim, _vmap_update_prim_attr, _bdim_at_any, _handle_broadcasting
+    _bdim_at_front, _vmap_clone_prim, _bdim_at_any, _handle_broadcasting
 
 
 @vmap_rules_getters.register(G.NLLLossGrad)
@@ -665,9 +665,8 @@ def get_grid_sampler_grad_vmap_rule(prim, axis_size):
 def get_upsample_grad_vmap_rule(prim, axis_size):
     """VmapRule for `UpsampleNearest3DGrad` and `UpsampleTrilinear3DGrad`."""
     cdhw_reverse_index = -4
-    input_size = prim.input_size
 
-    def vmap_rule(grad_bdim):
+    def vmap_rule(grad_bdim, isize_bdim, osize_bdim, scales_bdim):
         is_all_none, result = vmap_general_preprocess(prim, grad_bdim)
         if is_all_none:
             return result
@@ -679,11 +678,17 @@ def get_upsample_grad_vmap_rule(prim, axis_size):
         grad = F.reshape(grad, input_shape)
         real_in_shape = F.shape(grad)
 
+        isize, isize_dim = isize_bdim
+        osize, osize_dim = osize_bdim
+        scales, scales_dim = scales_bdim
+        if isize_dim is not None or osize_dim is not None or scales_dim is not None:
+            _raise_value_error(
+                "The source axis of `input_size`, `output_size` and `scales` must be None, but got {0}, {1} and {2}."
+                .format(isize_dim, osize_dim, scales_dim))
         # update batch dimension of input_size
-        new_input_size = (real_in_shape[0],) + input_size[1:]
-        _vmap_update_prim_attr(prim, 'input_size', new_input_size)
+        new_isize = (real_in_shape[0],) + isize[1:]
 
-        out = prim(grad)
+        out = prim(grad, new_isize, osize, scales)
         out_shape = F.shape(out)
         real_out_shape = grad_shape[:cdhw_reverse_index] + out_shape[cdhw_reverse_index:]
         out = F.reshape(out, real_out_shape)

mindspore/ops/_vmap/vmap_math_ops.py

@@ -66,6 +66,7 @@ def _broadcast_shape(nd, x_ndim, x_shape):
 @vmap_rules_getters.register(P.ApproximateEqual)
 @vmap_rules_getters.register(P.TruncateDiv)
 @vmap_rules_getters.register(P.TruncateMod)
+
 def get_broadcast_binary_op_vmap_rule(prim, axis_size):
     """VmapRule for binary operations with broadcasting, such as `Add` and `Sub`."""
 
@@ -216,8 +217,9 @@ def get_lerp_vamp_rule(prim, axis_size):
         # Both broadcast end and weight to start.
         else:
             weight_shape = F.shape(weight)
-            if (start_dim == end_dim and start_dim == weight_dim) and (
-                    start_shape == end_shape and start_shape == weight_shape):
+            is_dim_ok = start_dim == end_dim and start_dim == weight_dim
+            is_shape_ok = start_shape == end_shape and start_shape == weight_shape
+            if is_dim_ok and is_shape_ok:
                 out = prim(start, end, weight)
                 return out, start_dim
             start, end = broadcast_a_b_shape(start_bdim, end_bdim)
@@ -900,3 +902,4 @@ get_unop_vmap_rule = vmap_rules_getters.register(BesselK1e)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.Trunc)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.PopulationCount)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.Square)(get_unop_vmap_rule)
+get_unop_vmap_rule = vmap_rules_getters.register(P.Eps)(get_unop_vmap_rule)