mindspore 2.7.0__cp311-cp311-win_amd64.whl → 2.7.1__cp311-cp311-win_amd64.whl

mindspore/nn/probability/bijector/power_transform.py

@@ -13,8 +13,7 @@
 # limitations under the License.
 # ============================================================================
 """PowerTransform Bijector"""
-from mindspore.ops import operations as P
-from mindspore.ops import functional as F
+import mindspore.ops as ops
 from ..distribution._utils.utils import check_greater_equal_zero
 from ..distribution._utils.custom_ops import exp_generic, log_generic
 from .bijector import Bijector
@@ -76,16 +75,16 @@ class PowerTransform(Bijector):
         self._power = self._add_parameter(power, 'power')
         check_greater_equal_zero(self._power, 'Power')
 
-        self.pow = P.Pow()
-        self.dtypeop = P.DType()
-        self.cast = P.Cast()
-        self.equal_base = P.Equal()
+        self.pow = ops.Pow()
+        self.dtypeop = ops.DType()
+        self.cast = ops.Cast()
+        self.equal_base = ops.Equal()
         self.exp = exp_generic
-        self.expm1 = P.Expm1()
+        self.expm1 = ops.Expm1()
         self.log = log_generic
-        self.log1p = P.Log1p()
-        self.select_base = P.Select()
-        self.shape = P.Shape()
+        self.log1p = ops.Log1p()
+        self.select_base = ops.Select()
+        self.shape = ops.Shape()
 
     @property
     def power(self):
@@ -113,17 +112,17 @@ class PowerTransform(Bijector):
         power_local = self.cast_param_by_value(x, self.power)
 
         # broad cast the value of x and power
-        ones = F.fill(self.dtypeop(power_local), self.shape(x + power_local),
-                      1.)
+        ones = ops.fill(self.dtypeop(power_local), self.shape(x + power_local),
+                        1.)
         power_local = power_local * ones
         x = x * ones
         safe_power = self.select_base(
             self.equal_base(power_local,
-                            P.ZerosLike()(power_local)), ones, power_local)
+                            ops.ZerosLike()(power_local)), ones, power_local)
 
         forward_v = self.select_base(
             self.equal_base(power_local,
-                            P.ZerosLike()(power_local)), self.exp(x),
+                            ops.ZerosLike()(power_local)), self.exp(x),
             self.exp(self.log1p(x * safe_power) / safe_power))
         return forward_v
 
@@ -135,17 +134,17 @@ class PowerTransform(Bijector):
         power_local = self.cast_param_by_value(y, self.power)
 
         # broad cast the value of x and power
-        ones = F.fill(self.dtypeop(power_local), self.shape(y + power_local),
-                      1.)
+        ones = ops.fill(self.dtypeop(power_local), self.shape(y + power_local),
+                        1.)
         power_local = power_local * ones
         y = y * ones
         safe_power = self.select_base(
             self.equal_base(power_local,
-                            P.ZerosLike()(power_local)), ones, power_local)
+                            ops.ZerosLike()(power_local)), ones, power_local)
 
         inverse_v = self.select_base(
             self.equal_base(power_local,
-                            P.ZerosLike()(power_local)), self.log(y),
+                            ops.ZerosLike()(power_local)), self.log(y),
             self.expm1(self.log(y) * safe_power) / safe_power)
 
         return inverse_v
@@ -166,14 +165,14 @@ class PowerTransform(Bijector):
         power_local = self.cast_param_by_value(x, self.power)
 
         # broad cast the value of x and power
-        ones = F.fill(self.dtypeop(power_local), self.shape(x + power_local),
-                      1.)
+        ones = ops.fill(self.dtypeop(power_local), self.shape(x + power_local),
+                        1.)
         power_local = power_local * ones
         x = x * ones
 
         forward_log_j = self.select_base(
             self.equal_base(power_local,
-                            P.ZerosLike()(power_local)), x,
+                            ops.ZerosLike()(power_local)), x,
             (1. / power_local - 1) * self.log1p(x * power_local))
 
         return forward_log_j

mindspore/nn/probability/bijector/scalar_affine.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ============================================================================
 """Scalar Affine Bijector"""
-from mindspore.ops import operations as P
+import mindspore.ops as ops
 from ..distribution._utils.custom_ops import log_generic
 from .bijector import Bijector
 
@@ -86,10 +86,10 @@ class ScalarAffine(Bijector):
         self._scale = self._add_parameter(scale, 'scale')
         self._shift = self._add_parameter(shift, 'shift')
 
-        self.abs = P.Abs()
-        self.oneslike = P.OnesLike()
-        self.dtypeop = P.DType()
-        self.cast = P.Cast()
+        self.abs = ops.Abs()
+        self.oneslike = ops.OnesLike()
+        self.dtypeop = ops.DType()
+        self.cast = ops.Cast()
         self.log = log_generic
 
     @property

mindspore/nn/probability/bijector/softplus.py

@@ -14,8 +14,7 @@
 # ============================================================================
 """Softplus Bijector"""
 import numpy as np
-from mindspore.ops import operations as P
-from mindspore.ops import functional as F
+import mindspore.ops as ops
 from mindspore.nn.layer.activation import LogSigmoid
 from ..distribution._utils.custom_ops import exp_generic, log_generic
 from .bijector import Bijector
@@ -82,17 +81,17 @@ class Softplus(Bijector):
 
         self.exp = exp_generic
         self.log = log_generic
-        self.expm1 = P.Expm1()
-        self.abs = P.Abs()
-        self.dtypeop = P.DType()
-        self.cast = P.Cast()
-        self.greater = P.Greater()
-        self.less = P.Less()
+        self.expm1 = ops.Expm1()
+        self.abs = ops.Abs()
+        self.dtypeop = ops.DType()
+        self.cast = ops.Cast()
+        self.greater = ops.Greater()
+        self.less = ops.Less()
         self.log_sigmoid = LogSigmoid()
-        self.logicalor = P.LogicalOr()
-        self.select = P.Select()
-        self.shape = P.Shape()
-        self.sigmoid = P.Sigmoid()
+        self.logicalor = ops.LogicalOr()
+        self.select = ops.Select()
+        self.shape = ops.Shape()
+        self.sigmoid = ops.Sigmoid()
         self.softplus = self._softplus
         self.inverse_softplus = self._inverse_softplus
 
@@ -104,7 +103,7 @@ class Softplus(Bijector):
         too_large = self.greater(x, -self.threshold)
         too_small_value = self.exp(x)
         too_large_value = x
-        ones = F.fill(self.dtypeop(x), self.shape(x), 1.0)
+        ones = ops.fill(self.dtypeop(x), self.shape(x), 1.0)
         too_small_or_too_large = self.logicalor(too_small, too_large)
         x = self.select(too_small_or_too_large, ones, x)
         y = self.log(self.exp(x) + 1.0)
@@ -120,7 +119,7 @@ class Softplus(Bijector):
         too_large = self.greater(x, (-1) * self.threshold)
         too_small_value = self.log(x)
         too_large_value = x
-        ones = F.fill(self.dtypeop(x), self.shape(x), 1.0)
+        ones = ops.fill(self.dtypeop(x), self.shape(x), 1.0)
         too_small_or_too_large = self.logicalor(too_small, too_large)
         x = self.select(too_small_or_too_large, ones, x)
         y = x + self.log(self.abs(self.expm1((-1)*x)))

mindspore/nn/wrap/grad_reducer.py

@@ -140,34 +140,6 @@ def _tensors_allreduce_post(degree, mean, allreduce_filter, grad):
     return grad
 
 
-@reduce_opt.register("Tensor", "Bool", "Function", "Function", "Bool", "Tensor", "Bool")
-def _tensors_allreduce_ps(degree, mean, allgather, allreduce, allreduce_filter, grad, ps_parameter):
-    """
-    Apply allreduce on gradient.
-
-    Args:
-        degree (int): The mean coefficient.
-        mean (bool): When mean is true, the mean coefficient (degree) would apply on gradients.
-        allgather (Primitive): The communication operator for sparse gradients.
-        allreduce (Primitive): The communication operator for gradients.
-        allreduce_filter (bool): When it is true, allreduce would apply.
-        grad (Tensor): The gradient tensor before operation.
-        ps_parameter (bool): Use parameter server or not.
-
-    Returns:
-        Tensor, the gradient tensor after operation.
-    """
-    if ps_parameter:
-        return grad
-
-    if allreduce_filter:
-        grad = allreduce(grad)
-        if mean:
-            grad = ops.tensor_mul(grad, ops.cast(degree, ops.dtype(grad)))
-        return grad
-    return grad
-
-
 @reduce_opt.register("Tensor", "Bool", "Function", "Function", "Bool", "RowTensor")
 def _tensors_allreduce_with_sparse(degree, mean, allgather, allreduce, allreduce_filter, grad):
     """
@@ -193,37 +165,6 @@ def _tensors_allreduce_with_sparse(degree, mean, allgather, allreduce, allreduce
         grad = RowTensorInner(indices, dout, grad.dense_shape)
     return grad
 
-
-@reduce_opt.register("Tensor", "Bool", "Function", "Function", "Bool", "RowTensor", "Bool")
-def _tensors_allreduce_with_sparse_ps(degree, mean, allgather, allreduce, allreduce_filter, grad, ps_parameter):
-    """
-    Apply allgather on gradient instead of allreduce for sparse feature.
-    Allgather is a communication operation used for distributed deep learning.
-
-    Args:
-        degree (int): The mean coefficient.
-        mean (bool): When mean is true, the mean coefficient (degree) would apply on gradients.
-        allgather (Primitive): The communication operator for sparse gradients.
-        allreduce (Primitive): The communication operator for gradients.
-        allreduce_filter (bool): When it is true, allgather would apply.
-        grad (tuple): The indices, gradient tensor and tensor_shape before operation.
-        ps_parameter (bool): Use parameter server or not.
-
-    Returns:
-        RowTensor, the gradient after operation.
-    """
-    if ps_parameter:
-        return grad
-
-    if allreduce_filter:
-        indices = allgather(grad.indices)
-        dout = allgather(grad.values)
-        if mean:
-            dout = ops.tensor_mul(dout, ops.cast(degree, ops.dtype(dout)))
-        grad = RowTensorInner(indices, dout, grad.dense_shape)
-    return grad
-
-
 _get_datatype = ops.MultitypeFuncGraph("_get_datatype")
 
 
@@ -423,9 +364,6 @@ class DistributedGradReducer(Cell):
                 self.split_fusion = False
                 self.allreduce = AllReduce('sum', group).add_prim_attr('fusion', fusion_type)
         self.allgather = AllGather(group)
-        ps_filter = lambda x: x.is_param_ps
-        self.ps_parameters = tuple(ps_filter(x) for x in parameters)
-        self.enable_parameter_server = any(self.ps_parameters)
         self.mode = context.get_context("mode")
         self.enable_tuple_broaden = True
 
@@ -446,19 +384,11 @@ class DistributedGradReducer(Cell):
         grads = self.map_(ops.partial(_cast_datatype, mstype.float32), grads)
 
         if self.split_fusion:
-            if self.enable_parameter_server:
-                new_grad = self.map_(ops.partial(reduce_opt, self.degree, self.mean, self.allgather),
-                                     self.op_list, self.allreduce_filter, grads, self.ps_parameters)
-            else:
-                new_grad = self.map_(ops.partial(reduce_opt, self.degree, self.mean, self.allgather),
-                                     self.op_list, self.allreduce_filter, grads)
+            new_grad = self.map_(ops.partial(reduce_opt, self.degree, self.mean, self.allgather),
+                                 self.op_list, self.allreduce_filter, grads)
         else:
-            if self.enable_parameter_server:
-                new_grad = self.map_(ops.partial(reduce_opt, self.degree, self.mean, self.allgather,
-                                                 self.allreduce), self.allreduce_filter, grads, self.ps_parameters)
-            else:
-                new_grad = self.map_(ops.partial(reduce_opt, self.degree, self.mean, self.allgather,
-                                                 self.allreduce), self.allreduce_filter, grads)
+            new_grad = self.map_(ops.partial(reduce_opt, self.degree, self.mean, self.allgather,
+                                             self.allreduce), self.allreduce_filter, grads)
         new_grad = self.map_(ops.partial(_cast_datatype), datatypes, new_grad)
         return new_grad
 

mindspore/numpy/array_creations.py

@@ -2622,7 +2622,7 @@ def pad(arr, pad_width, mode="constant", stat_length=None, constant_values=0,
             unique pad widths for each axis. ``((before, after),)`` yields same
             before and after pad for each axis. ``(pad,)`` or int is a shortcut
             for ``before = after = pad width`` for all axes.
-        mode (string, optional):
+        mode (str, optional):
             One of the following string values:
 
             - constant (default): Pads with a constant value.
@@ -2660,7 +2660,7 @@ def pad(arr, pad_width, mode="constant", stat_length=None, constant_values=0,
             unique end values for each axis. ``((before, after),)`` yields same before
             and after end values for each axis. ``(constant,)`` or ``constant``
             is a shortcut for ``before = after = constant`` for all axes. Default: ``0`` .
-        reflect_type(string, optional) can choose between \'even\' and \'odd\'. Used in
+        reflect_type(str, optional) can choose between \'even\' and \'odd\'. Used in
             \'reflect\', and \'symmetric\'. The \'even\' style is the default with an
             unaltered reflection around the edge value. For the \'odd\' style, the extended
             part of the `arr` is created by subtracting the reflected values from two times

mindspore/numpy/fft.py

@@ -185,7 +185,7 @@ def rfft(a, n=None, axis=-1, norm=None):
             Default: ``None``.
         axis (int, optional): Axis over which to compute the `rfft`.
             Default: ``-1``, which means the last axis of `a` is used.
-        norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"``.
+        norm (str, optional): Normalization mode. Default: ``None`` that means ``"backward"``.
             Three modes are defined as,
 
             - ``"backward"`` (no normalization).
@@ -224,7 +224,7 @@ def irfft(a, n=None, axis=-1, norm=None):
             Default: ``None``.
         axis (int, optional): Axis over which to compute the `irfft`.
             Default: ``-1``, which means the last axis of `a` is used.
-        norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"``.
+        norm (str, optional): Normalization mode. Default: ``None`` that means ``"backward"``.
             Three modes are defined as,
 
             - ``"backward"`` (normalize by :math:`1/n`).
@@ -266,7 +266,7 @@ def fft2(a, s=None, axes=(-2, -1), norm=None):
             Default: ``None`` , which does not need to process `a`.
         axes (tuple[int], optional): The dimension along which to take the one dimensional `fft2`.
             Default: ``(-2, -1)`` , which means transform the last two dimension of `a`.
-        norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
+        norm (str, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
             Three modes are defined as, where :math: `n = prod(s)`
 
             - ``"backward"`` (no normalization).
@@ -361,7 +361,7 @@ def fftn(a, s=None, axes=None, norm=None):
         axes (tuple[int], optional): The dimension along which to take the one dimensional `fftn`.
             Default: ``None`` , which means transform the all dimension of `a`,
             or the last `len(s)` dimensions if s is given.
-        norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
+        norm (str, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
             Three modes are defined as, where :math: `n = prod(s)`
 
             - ``"backward"`` (no normalization).
@@ -409,7 +409,7 @@ def ifftn(a, s=None, axes=None, norm=None):
         axes (tuple[int], optional): The dimension along which to take the one dimensional `ifftn`.
             Default: ``None`` , which means transform the all dimension of `a`,
             or the last `len(s)` dimensions if s is given.
-        norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
+        norm (str, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
             Three modes are defined as, where :math: `n = prod(s)`
 
             - ``"backward"`` (normalize by :math:`1/n`).
@@ -457,7 +457,7 @@ def rfft2(a, s=None, axes=(-2, -1), norm=None):
             Default: ``None`` , which does not need to process `a`.
         axes (tuple[int], optional): The dimension along which to take the one dimensional `rfft2`.
             Default: ``(-2, -1)`` , which means transform the last two dimension of `a`.
-        norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
+        norm (str, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
             Three modes are defined as, where :math: `n = prod(s)`
 
             - ``"backward"`` (no normalization).
@@ -502,7 +502,7 @@ def irfft2(a, s=None, axes=(-2, -1), norm=None):
             Default: ``None`` , the axes[-1] of the `a` will be zero-padded to :math:`2*(a.shape[axes[-1]]-1)`.
         axes (tuple[int], optional): The dimension along which to take the one dimensional `irfft2`.
             Default: ``(-2, -1)`` , which means transform the last two dimension of `a`.
-        norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
+        norm (str, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
             Three modes are defined as, where :math: `n = prod(s)`
 
             - ``"backward"`` (normalize by :math:`1/n`).
@@ -551,7 +551,7 @@ def rfftn(a, s=None, axes=None, norm=None):
         axes (tuple[int], optional): The dimension along which to take the one dimensional `rfftn`.
             Default: ``None`` , which means transform the all dimension of `a`,
             or the last `len(s)` dimensions if s is given.
-        norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
+        norm (str, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
             Three modes are defined as, where :math: `n = prod(s)`
 
             - ``"backward"`` (no normalization).
@@ -599,7 +599,7 @@ def irfftn(a, s=None, axes=None, norm=None):
         axes (tuple[int], optional): The dimension along which to take the one dimensional `irfftn`.
             Default: ``None`` , which means transform the all dimension of `a`,
             or the last `len(s)` dimensions if s is given.
-        norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
+        norm (str, optional): Normalization mode. Default: ``None`` that means ``"backward"`` .
             Three modes are defined as, where :math: `n = prod(s)`
 
             - ``"backward"`` (normalize by :math:`1/n`).

mindspore/{nn/reinforcement → onnx}/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2021 Huawei Technologies Co., Ltd
+# 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.
@@ -12,13 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""
-TensorArray.
-"""
+"""onnx module."""
+
 from __future__ import absolute_import
 
-from mindspore.nn.reinforcement.tensor_array import TensorArray
+from .onnx_export import export
 
-__all__ = [
-    "TensorArray",
-]
+__all__ = ["export"]

mindspore/onnx/onnx_export.py

@@ -0,0 +1,137 @@
+# 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.
+# ============================================================================
+
+"""Model export to ONNX."""
+from __future__ import absolute_import
+from __future__ import division
+
+import os
+
+import mindspore.nn as nn
+from mindspore import log as logger
+from mindspore._checkparam import check_input_dataset
+from mindspore import _checkparam as Validator
+from mindspore.common.api import _cell_graph_executor as _executor
+from mindspore.train.serialization import _calculation_net_size
+from mindspore.dataset.engine.datasets import Dataset
+
+PROTO_LIMIT_SIZE = 1024 * 1024 * 2
+
+
+def export(net, *inputs, file_name, input_names=None, output_names=None, export_params=True,
+           keep_initializers_as_inputs=False, dynamic_axes=None):
+    """
+    Export the MindSpore network into an ONNX model.
+
+    Note:
+        - Support exporting network larger than 2GB. When the network exceeds 2GB,
+          parameters are saved in additional binary files stored in the same directory as the ONNX file.
+        - When `file_name` does not have a suffix, the system will automatically add the suffix `.onnx` .
+
+    Args:
+        net (Union[Cell, function]): MindSpore network.
+        inputs (Union[Tensor, list, tuple, Number, bool]): It represents the inputs of the `net` , if the network has
+            multiple inputs, set them together.
+        file_name (str): File name of the model to be exported.
+        input_names (list, optional): Names to assign to the input nodes of the graph, in order. Default: ``None`` .
+        output_names (list, optional): Names to assign to the output nodes of the graph, in order. Default: ``None`` .
+        export_params (bool, optional): If false, parameters (weights) will not be exported,
+            parameters will add input nodes as input of the graph. Default: ``True`` .
+        keep_initializers_as_inputs (bool, optional): If True, all the initializers (model parameters/weights) will
+            add as inputs to the graph. This allows modifying any or all weights when running the exported ONNX model.
+            Default: ``False`` .
+        dynamic_axes (dict[str, dict[int, str]], optional): To specify axes of input tensors as dynamic (at runtime).
+            Default: ``None`` .
+
+            - Set a dict with scheme: {input_node_name: {axis_index:axis_name}},
+              for example, {"input1": {0:"batch_size", 1: "seq_len"}, "input2": {{0:"batch_size"}}.
+            - By default, the shapes of all input tensors in the exported model exactly match those specified in
+              `inputs`.
+
+    Raises:
+        ValueError: If the parameter `net` is not :class:`mindspore.nn.Cell`.
+        ValueError: If the parameter `input_names` is not list type.
+        ValueError: If the parameter `output_names` is not list type
+        ValueError: If the parameter `dynamic_axes` is not dict type.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> from mindspore import Tensor
+        >>>
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> input_tensor = Tensor(np.ones([1, 1, 32, 32]).astype(np.float32))
+        >>> ms.onnx.export(net, input_tensor, file_name='lenet.onnx', input_names=['input1'], output_names=['output1'])
+
+    """
+    Validator.check_file_name_by_regular(file_name)
+    logger.info("exporting model file:%s format:%s.", file_name, "ONNX")
+    Validator.check_isinstance("net", net, nn.Cell)
+    input_names = input_names or []
+    Validator.check_isinstance("input_names", input_names, list)
+    output_names = output_names or []
+    Validator.check_isinstance("output_names", output_names, list)
+    dynamic_axes = dynamic_axes or {}
+    Validator.check_isinstance("dynamic_axes", dynamic_axes, dict)
+
+    if check_input_dataset(*inputs, dataset_type=Dataset):
+        raise ValueError(f"Can not support dataset as inputs to export ONNX model.")
+
+    cell_mode = net.training
+    net.set_train(mode=False)
+
+    extra_save_params = False
+    total_size = _calculation_net_size(net)
+    if total_size > PROTO_LIMIT_SIZE:
+        logger.warning('Network size is: {}G, it exceeded the protobuf: {}G limit, now parameters in network are saved '
+                       'in external data files.'.format(total_size / 1024 / 1024, PROTO_LIMIT_SIZE / 1024 / 1024))
+        extra_save_params = True
+
+    phase_name = 'export.onnx'
+    graph_id, _ = _executor.compile(net, *inputs, phase=phase_name, do_convert=False)
+
+    abs_file_name = os.path.abspath(file_name)
+    if not abs_file_name.endswith('.onnx'):
+        abs_file_name += ".onnx"
+
+    dir_path = os.path.dirname(abs_file_name)
+    if dir_path and not os.path.exists(dir_path):
+        os.makedirs(dir_path, exist_ok=True)
+
+    abs_file_dir = os.path.dirname(abs_file_name) if extra_save_params else ""
+
+    onnx_stream = _executor._get_onnx_func_graph_proto(obj=net, exec_id=graph_id, input_names=input_names,
+                                                       output_names=output_names, export_params=export_params,
+                                                       keep_initializers_as_inputs=keep_initializers_as_inputs,
+                                                       dynamic_axes=dynamic_axes, extra_save_params=extra_save_params,
+                                                       save_file_dir=abs_file_dir)
+    if onnx_stream is None:
+        raise RuntimeError("Export onnx model failed, ensure that the model has been compiled correctly")
+
+    try:
+        with open(abs_file_name, 'wb') as f:
+            f.write(onnx_stream)
+
+        if os.path.getsize(abs_file_name) != len(onnx_stream):
+            logger.warning("ONNX file size doesn't match expected value, but proceeding continue.")
+
+    except IOError as e:
+        logger.error(f"Failed to write ONNX file: {e}")
+        if os.path.exists(abs_file_name):
+            os.remove(abs_file_name)
+
+    net.set_train(mode=cell_mode)

mindspore/ops/__init__.py

@@ -37,6 +37,7 @@ from mindspore.ops.functional_overload import all_gather_matmul, matmul_reduce_s
 from mindspore.ops.composite import *
 from mindspore.ops.operations import *
 from mindspore.ops.function import *
+from mindspore.ops.communication import *
 from mindspore.ops.functional import *
 from mindspore.ops._utils import arg_dtype_cast, arg_handler
 
@@ -55,4 +56,5 @@ __all__.extend(composite.__all__)
 __all__.extend(operations.__all__)
 __all__.extend(functional.__all__)
 __all__.extend(function.__all__)
+__all__.extend(communication.__all__)
 __all__.extend(auto_generate.__all__)