mindspore 2.3.0__cp39-none-any.whl → 2.3.0rc2__cp39-none-any.whl

mindspore/ops/auto_generate/pyboost_inner_prim.py

@@ -19,14 +19,10 @@ from mindspore._c_expression import ArgMaxWithValuePrim_
 from mindspore._c_expression import ArgMinWithValuePrim_
 from mindspore._c_expression import BatchMatMulPrim_
 from mindspore._c_expression import BatchNormGradExtPrim_
-from mindspore._c_expression import BinaryCrossEntropyGradPrim_
-from mindspore._c_expression import BinaryCrossEntropyPrim_
-from mindspore._c_expression import BCEWithLogitsLossPrim_
 from mindspore._c_expression import BroadcastToPrim_
 from mindspore._c_expression import ConcatPrim_
 from mindspore._c_expression import ConvolutionGradPrim_
 from mindspore._c_expression import ConvolutionPrim_
-from mindspore._c_expression import EluExtPrim_
 from mindspore._c_expression import FFNExtPrim_
 from mindspore._c_expression import FlashAttentionScoreGradPrim_
 from mindspore._c_expression import FlashAttentionScorePrim_
@@ -34,25 +30,22 @@ from mindspore._c_expression import GridSampler2DGradPrim_
 from mindspore._c_expression import GridSampler2DPrim_
 from mindspore._c_expression import GridSampler3DGradPrim_
 from mindspore._c_expression import GridSampler3DPrim_
-from mindspore._c_expression import IsClosePrim_
 from mindspore._c_expression import MatMulPrim_
 from mindspore._c_expression import MaxPoolGradWithIndicesPrim_
 from mindspore._c_expression import MaxPoolGradWithMaskPrim_
 from mindspore._c_expression import MaxPoolWithIndicesPrim_
 from mindspore._c_expression import MaxPoolWithMaskPrim_
 from mindspore._c_expression import OneHotExtPrim_
+from mindspore._c_expression import QuantBatchMatmulPrim_
 from mindspore._c_expression import ReduceAllPrim_
 from mindspore._c_expression import ReduceAnyPrim_
 from mindspore._c_expression import ReverseV2Prim_
-from mindspore._c_expression import RmsNormPrim_
-from mindspore._c_expression import SearchSortedPrim_
 from mindspore._c_expression import SoftmaxPrim_
 from mindspore._c_expression import StackExtPrim_
+from mindspore._c_expression import TrilPrim_
 from mindspore._c_expression import TriuPrim_
 from mindspore._c_expression import UpsampleTrilinear3DGradPrim_
 from mindspore._c_expression import UpsampleTrilinear3DPrim_
-from mindspore._c_expression import GroupedMatmulPrim_
-from mindspore._c_expression import QuantBatchMatmulPrim_
 from mindspore._c_expression import WeightQuantBatchMatmulPrim_
 
 
@@ -92,33 +85,6 @@ class _PyboostBatchNormGradExtPrim(BatchNormGradExtPrim_):
 batch_norm_grad_ext_impl = _PyboostBatchNormGradExtPrim()
 
 
-class _PyboostBinaryCrossEntropyGradPrim(BinaryCrossEntropyGradPrim_):
-    def __call__(self, input, target, grad_output, weight, reduction):
-        converted_reduction = str_to_enum(reduction)
-        return _convert_stub(super().__call__(input, target, grad_output, weight, reduction))
-
-
-binary_cross_entropy_grad_impl = _PyboostBinaryCrossEntropyGradPrim()
-
-
-class _PyboostBinaryCrossEntropyPrim(BinaryCrossEntropyPrim_):
-    def __call__(self, input, target, weight, reduction):
-        converted_reduction = str_to_enum(reduction)
-        return _convert_stub(super().__call__(input, target, weight, reduction))
-
-
-binary_cross_entropy_impl = _PyboostBinaryCrossEntropyPrim()
-
-
-class _PyboostBCEWithLogitsLossPrim(BCEWithLogitsLossPrim_):
-    def __call__(self, input, target, weight, posWeight, reduction):
-        converted_reduction = str_to_enum(reduction)
-        return _convert_stub(super().__call__(input, target, weight, posWeight, reduction))
-
-
-binary_cross_entropy_with_logits_impl = _PyboostBCEWithLogitsLossPrim()
-
-
 class _PyboostBroadcastToPrim(BroadcastToPrim_):
     def __call__(self, input, shape):
 
@@ -161,15 +127,6 @@ class _PyboostConvolutionPrim(ConvolutionPrim_):
 convolution_impl = _PyboostConvolutionPrim()
 
 
-class _PyboostEluExtPrim(EluExtPrim_):
-    def __call__(self, input, alpha):
-
-        return _convert_stub(super().__call__(input, alpha))
-
-
-elu_ext_impl = _PyboostEluExtPrim()
-
-
 class _PyboostFFNExtPrim(FFNExtPrim_):
     def __call__(self, x, weight1, weight2, expertTokens, bias1, bias2, scale, offset, deqScale1, deqScale2, antiquant_scale1, antiquant_scale2, antiquant_offset1, antiquant_offset2, activation, inner_precise):
         converted_activation = str_to_enum(activation)
@@ -237,15 +194,6 @@ class _PyboostGridSampler3DPrim(GridSampler3DPrim_):
 grid_sampler_3d_impl = _PyboostGridSampler3DPrim()
 
 
-class _PyboostIsClosePrim(IsClosePrim_):
-    def __call__(self, input, other, rtol, atol, equal_nan):
-
-        return _convert_stub(super().__call__(input, other, rtol, atol, equal_nan))
-
-
-isclose_impl = _PyboostIsClosePrim()
-
-
 class _PyboostMatMulPrim(MatMulPrim_):
     def __call__(self, input, mat2, transpose_a, transpose_b):
 
@@ -312,6 +260,15 @@ class _PyboostOneHotExtPrim(OneHotExtPrim_):
 one_hot_ext_impl = _PyboostOneHotExtPrim()
 
 
+class _PyboostQuantBatchMatmulPrim(QuantBatchMatmulPrim_):
+    def __call__(self, x1, x2, scale, offset, bias, transpose_x1, transpose_x2, dtype):
+
+        return _convert_stub(super().__call__(x1, x2, scale, offset, bias, transpose_x1, transpose_x2, dtype))
+
+
+quant_batch_matmul_impl = _PyboostQuantBatchMatmulPrim()
+
+
 class _PyboostReduceAllPrim(ReduceAllPrim_):
     def __call__(self, input, axis, keep_dims):
 
@@ -339,24 +296,6 @@ class _PyboostReverseV2Prim(ReverseV2Prim_):
 reverse_v2_impl = _PyboostReverseV2Prim()
 
 
-class _PyboostRmsNormPrim(RmsNormPrim_):
-    def __call__(self, x, gamma, epsilon):
-
-        return _convert_stub(super().__call__(x, gamma, epsilon))
-
-
-rms_norm_impl = _PyboostRmsNormPrim()
-
-
-class _PyboostSearchSortedPrim(SearchSortedPrim_):
-    def __call__(self, sorted_sequence, values, sorter, dtype, right):
-
-        return _convert_stub(super().__call__(sorted_sequence, values, sorter, dtype, right))
-
-
-searchsorted_impl = _PyboostSearchSortedPrim()
-
-
 class _PyboostSoftmaxPrim(SoftmaxPrim_):
     def __call__(self, input, axis):
 
@@ -375,6 +314,15 @@ class _PyboostStackExtPrim(StackExtPrim_):
 stack_ext_impl = _PyboostStackExtPrim()
 
 
+class _PyboostTrilPrim(TrilPrim_):
+    def __call__(self, input, diagonal):
+
+        return _convert_stub(super().__call__(input, diagonal))
+
+
+tril_impl = _PyboostTrilPrim()
+
+
 class _PyboostTriuPrim(TriuPrim_):
     def __call__(self, input, diagonal):
 
@@ -402,24 +350,6 @@ class _PyboostUpsampleTrilinear3DPrim(UpsampleTrilinear3DPrim_):
 upsample_trilinear3d_impl = _PyboostUpsampleTrilinear3DPrim()
 
 
-class _PyboostGroupedMatmulPrim(GroupedMatmulPrim_):
-    def __call__(self, x, weight, bias, scale, offset, antiquant_scale, antiquant_offset, group_list, split_item, group_type):
-
-        return _convert_stub(super().__call__(x, weight, bias, scale, offset, antiquant_scale, antiquant_offset, group_list, split_item, group_type))
-
-
-grouped_matmul_impl = _PyboostGroupedMatmulPrim()
-
-
-class _PyboostQuantBatchMatmulPrim(QuantBatchMatmulPrim_):
-    def __call__(self, x1, x2, scale, offset, bias, transpose_x1, transpose_x2, dtype):
-
-        return _convert_stub(super().__call__(x1, x2, scale, offset, bias, transpose_x1, transpose_x2, dtype))
-
-
-quant_batch_matmul_impl = _PyboostQuantBatchMatmulPrim()
-
-
 class _PyboostWeightQuantBatchMatmulPrim(WeightQuantBatchMatmulPrim_):
     def __call__(self, x, weight, antiquant_scale, antiquant_offset, quant_scale, quant_offset, bias, transpose_x, transpose_weight, antiquant_group_size):
 

mindspore/ops/composite/base.py

@@ -414,8 +414,9 @@ class GradOperation(GradOperation_):
         else:
             # Check if fn have run already
             if not _pynative_executor.check_run(grad, fn, weights, None, *args, **new_kwargs):
-                _pynative_executor.set_grad_flag(True)
+                fn.set_grad()
                 fn(*args, **new_kwargs)
+                fn.set_grad(False)
 
 
 class _TaylorOperation(TaylorOperation_):
@@ -653,8 +654,9 @@ class _Grad(GradOperation_):
         else:
             # Check if fn has run already.
             if not _pynative_executor.check_run(grad, fn, weights, self.grad_position, *args, **new_kwargs):
-                _pynative_executor.set_grad_flag(True)
+                fn.set_grad()
                 outputs = fn(*args, **new_kwargs)
+                fn.set_grad(False)
                 return outputs
         if (self.get_value or self.has_aux) and not outputs:
             outputs = fn(*args, **new_kwargs)
@@ -713,8 +715,9 @@ class MultitypeFuncGraph(MultitypeFuncGraph_):
         >>> # `add` is a metagraph object which will add two objects according to
         >>> # input type using ".register" decorator.
         >>> from mindspore import Tensor
-        >>> from mindspore import dtype as mstype
         >>> from mindspore import ops
+        >>> from mindspore import dtype as mstype
+        >>> import mindspore.ops as ops
         >>>
         >>> tensor_add = ops.Add()
         >>> add = ops.MultitypeFuncGraph('add')

mindspore/ops/composite/math_ops.py

@@ -112,7 +112,7 @@ def mm(input, mat2):
 
     Examples:
         >>> import mindspore as ms
-        >>> from mindspore import ops
+        >>> import mindspore.ops as ops
         >>> import numpy as np
         >>> x1 = ms.Tensor(np.random.rand(2, 3), ms.float32)
         >>> x2 = ms.Tensor(np.random.rand(3, 4), ms.float32)

mindspore/ops/composite/multitype_ops/_compile_utils.py

@@ -218,8 +218,8 @@ def _tensor_setitem(self, index, value):
     return output
 
 
-setattr(tensor_operator_registry, "__getitem__", _tensor_getitem)
-setattr(tensor_operator_registry, "__setitem__", _tensor_setitem)
+tensor_operator_registry.register("__getitem__", _tensor_getitem)
+tensor_operator_registry.register("__setitem__", _tensor_setitem)
 
 
 def _tensor_add(self, other):
@@ -288,15 +288,15 @@ def _tensor_floordiv(self, other):
     return F.floordiv(self, other)
 
 
-setattr(tensor_operator_registry, '__add__', _tensor_add)
-setattr(tensor_operator_registry, '__sub__', _tensor_sub)
-setattr(tensor_operator_registry, '__mul__', _tensor_mul)
-setattr(tensor_operator_registry, '__matmul__', _tensor_matmul)
-setattr(tensor_operator_registry, '__truediv__', _tensor_div)
-setattr(tensor_operator_registry, '__mod__', _tensor_mod)
-setattr(tensor_operator_registry, '__pow__', _tensor_pow)
-setattr(tensor_operator_registry, '__rpow__', _tensor_rpow)
-setattr(tensor_operator_registry, '__floordiv__', _tensor_floordiv)
+tensor_operator_registry.register('__add__', _tensor_add)
+tensor_operator_registry.register('__sub__', _tensor_sub)
+tensor_operator_registry.register('__mul__', _tensor_mul)
+tensor_operator_registry.register('__matmul__', _tensor_matmul)
+tensor_operator_registry.register('__truediv__', _tensor_div)
+tensor_operator_registry.register('__mod__', _tensor_mod)
+tensor_operator_registry.register('__pow__', _tensor_pow)
+tensor_operator_registry.register('__rpow__', _tensor_rpow)
+tensor_operator_registry.register('__floordiv__', _tensor_floordiv)
 
 
 def _scalar_to_tensor(input_x):
@@ -356,8 +356,8 @@ def tensor_itemset(data, *args):
     return tensor_itemset_with_number(data, args[0])
 
 
-setattr(tensor_operator_registry, "item", tensor_item)
-setattr(tensor_operator_registry, "itemset", tensor_itemset)
+tensor_operator_registry.register("item", tensor_item)
+tensor_operator_registry.register("itemset", tensor_itemset)
 
 
 def tensor_itemset_with_number(data, number_value):
@@ -1204,11 +1204,7 @@ def tensor_setitem_by_ellipsis_with_tensor(data, value):
     value = value.astype(data_dtype)
 
     value_shape = F.shape(value)
-
-    if len(value_shape) > len(data_shape):
-        source_shape = data_shape
-    else:
-        source_shape = value_shape
+    source_shape = const_utils.get_source_shape(data_shape, value_shape)
     value = F.reshape(value, source_shape)
     data = F.broadcast_to(value, data_shape)
     return data
@@ -1233,10 +1229,7 @@ def tensor_setitem_by_bool(data, index, value):
 
     if index:
         value_shape = F.shape(value)
-        if len(value_shape) > len(data_shape):
-            source_shape = data_shape
-        else:
-            source_shape = value_shape
+        source_shape = const_utils.get_source_shape(data_shape, value_shape)
         value = F.reshape(value, source_shape)
         data = F.broadcast_to(value, data_shape)
     return data
@@ -1403,7 +1396,7 @@ def reduce_(a, reduce_fn, cmp_fn=None, axis=None, keepdims=False, initial=None,
     return reduce_fn(a, axes).astype(dtype)
 
 
-setattr(tensor_operator_registry, "reduce", reduce_)
+tensor_operator_registry.register("reduce", reduce_)
 
 
 def check_indices(dims, indices, mode, allow_negative_index=True):
@@ -1430,7 +1423,7 @@ def check_indices(dims, indices, mode, allow_negative_index=True):
     return clipped
 
 
-setattr(tensor_operator_registry, 'check_indices', check_indices)
+tensor_operator_registry.register('check_indices', check_indices)
 
 
 def convert_slice_to_tensor(index, final_shape, slice_cnt, broadcast_shape, slice_shapes, fancy_position):

mindspore/ops/composite/multitype_ops/_constexpr_utils.py

@@ -216,7 +216,7 @@ def make_tensor(a, dtype=mstype.int64, data_shape=None, dim_size=None):
 
     return Tensor(a, dtype)
 
-setattr(tensor_operator_registry, 'make_tensor', make_tensor)
+tensor_operator_registry.register('make_tensor', make_tensor)
 
 
 def judge_data_dim(data_dim, min_data_dim=0, max_data_dim=8):

mindspore/ops/extend/__init__.py

@@ -33,7 +33,7 @@ from . import (
     nn_func,
 )
 
-from .array_func import gather, max, min, one_hot
+from .array_func import gather, max, min, one_hot, narrow
 from .math_func import (
     baddbmm,
     bmm,
@@ -44,7 +44,8 @@ from .math_func import (
 from .nn_func import (
     conv2d,
     max_pool2d,
-    leaky_relu_ext
+    leaky_relu_ext,
+    batch_norm
 )
 
 __all__ = []

mindspore/ops/extend/array_func.py

@@ -21,12 +21,52 @@ Array Operators
 from mindspore.common import Tensor
 from mindspore.ops.operations.array_ops import ArgMaxWithValue, ArgMinWithValue
 from mindspore.ops._primitive_cache import _get_cache_prim
-from mindspore.ops.auto_generate.gen_ops_prim import gather_d_op
+from mindspore.ops.auto_generate.gen_ops_prim import gather_d_op, slice_ext_op, OneHotExt
 from mindspore.ops.auto_generate.gen_ops_def import max_, min_
-from mindspore.ops.auto_generate.pyboost_inner_prim import _PyboostOneHotExtPrim
-one_hot_ext_impl = _PyboostOneHotExtPrim()
+from mindspore import _checkparam as validator
+
 
 # define Primitive global variables
+def narrow(input, dim, start, length):
+    """
+    Returns a narrowed tensor from input tensor, and
+    the dimension axis is input from start to start + length.
+
+    Args:
+        input (Tensor): the tensor to narrow.
+        dim (int): dimension  along which to narrow.
+        start (int): the starting dimension.
+        length (int): the distance to the ending dimension.
+
+    Returns:
+        Tensor.
+
+        - output (Tensors) - The narrowed tensor.
+
+    Raises:
+        TypeError: If the input is not a tensor or tuple or list of tensors.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import ops
+        >>> from mindspore import Tensor
+        >>> x = Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], mindspore.int32)
+        >>> output = ops.narrow(x, 0, 0, 2)
+        >>> print(output)
+        [[ 1 2 3]
+         [ 4 5 6]]
+        >>> output = ops.narrow(x, 1, 1, 2)
+        >>> print(output)
+        [[ 2 3]
+         [ 5 6]
+         [ 8 9]]
+    """
+    validator.check_value_type("input", input, Tensor, "narrow")
+    return slice_ext_op(input, dim, start, start+length, 1)
+
 
 def gather(input, dim, index):
     r"""
@@ -48,7 +88,7 @@ def gather(input, dim, index):
         index (Tensor): The index tensor, with int32 or int64 data type. An valid `index` should be:
 
             - `index.rank == input.rank`;
-            - for `axis != dim`, `index.shape[axis] <= input.shape[axis]`;
+            - `index.shape[axis] <= input.shape[axis]` where axis goes through all dimensions of `input` except `dim`;
             - the value of `index` is in range `[-input.shape[dim], input.shape[dim])`.
 
     Returns:
@@ -72,7 +112,7 @@ def gather(input, dim, index):
         >>> output = ops.extend.gather(input, 1, index)
         >>> print(output)
         [[-0.1 -0.1]
-         [0.5   0.5]]
+        [ 0.5  0.5]]
     """
     return gather_d_op(input, dim, index)
 
@@ -92,7 +132,7 @@ def max(input, dim=None, keepdim=False):
         and same dtype as `input`.
 
         tuple (Tensor) if `dim` is not the default value ``None`` , tuple of 2 tensors, containing the maximum
-        value of the input tensor along the given dimension `dim` and the corresponding index.
+        value of the input tensor along the given dimension `dim` and the corresponding index:
 
         - **values (Tensor)** - The maximum value of input tensor along the given dimension `dim`, with same dtype as
           `input`. If `keepdim` is ``True`` , the shape of output tensors is :math:`(input_1, input_2, ...,
@@ -141,7 +181,7 @@ def min(input, dim=None, keepdim=False):
         and same dtype as `input`.
 
         tuple (Tensor) if `dim` is not the default value ``None`` , tuple of 2 tensors, containing the minimum value
-        of the input tensor along the given dimension `dim` and the corresponding index.
+        of the input tensor along the given dimension `dim` and the corresponding index:
 
         - **values (Tensor)** - The minimum value of input tensor along the given dimension `dim`, with same dtype as
           `input`. If `keepdim` is ``True`` , the shape of output tensors is :math:`(input_1, input_2, ...,
@@ -184,7 +224,7 @@ def one_hot(tensor, num_classes):
     Args:
         tensor (Tensor): A tensor of indices. Tensor of shape :math:`(X_0, \ldots, X_n)`.
             Data type must be int32 or int64.
-        num_classes (int): A scalar defining the depth of the one-hot dimension.
+        num_classes (Union[int, Tensor]): A scalar defining the depth of the one-hot dimension.
 
     Returns:
         Tensor, one-hot tensor.
@@ -200,7 +240,7 @@ def one_hot(tensor, num_classes):
     Examples:
         >>> import mindspore
         >>> import numpy as np
-        >>> from mindspore import ops
+        >>> import mindspore.ops as ops
         >>> from mindspore import Tensor
         >>> tensor = Tensor(np.array([0, 1, 2]), mindspore.int32)
         >>> num_classes = 3
@@ -212,7 +252,8 @@ def one_hot(tensor, num_classes):
     """
     on_value = Tensor(1, dtype=tensor.dtype)
     off_value = Tensor(0, dtype=tensor.dtype)
-    return one_hot_ext_impl(tensor, num_classes, on_value, off_value, -1)
+    onehot = _get_cache_prim(OneHotExt)(-1)
+    return onehot(tensor, num_classes, on_value, off_value)
 
 
 __all__ = ['gather', 'max', 'min', 'one_hot']

mindspore/ops/extend/nn_func.py

@@ -21,6 +21,8 @@ NN Operators with better performance
 from mindspore.ops._primitive_cache import _get_cache_prim
 from mindspore.ops.auto_generate.gen_ops_prim import Convolution, ConstantPadND, MaxPoolWithIndices, MaxPoolWithMask
 from mindspore.ops.auto_generate import leaky_relu_ext
+from mindspore.ops.auto_generate import BatchNormExt
+from mindspore import ops
 from mindspore import _checkparam as validator
 
 
@@ -232,7 +234,7 @@ def max_pool2d(input, kernel_size, stride=None, padding=0, dilation=1, *, ceil_m
         \text{input}(N_i, C_j, s_0 \times h + m, s_1 \times w + n)
 
     .. warning::
-        Only support on Atlas A2 training series.
+        Only support on Atlas training series.
 
     Args:
         input (Tensor): Tensor of shape :math:`(N_{in}, C_{in}, H_{in}, W_{in})` with data type of float32
@@ -305,4 +307,78 @@ def max_pool2d(input, kernel_size, stride=None, padding=0, dilation=1, *, ceil_m
     return out
 
 
-__all__ = ['conv2d', 'max_pool2d', 'leaky_relu_ext']
+def batch_norm(input, running_mean, running_var, weight=None, bias=None, training=False, momentum=0.1, eps=1e-5):
+    r"""
+    Batch Normalization for input data and updated parameters.
+
+    Batch Normalization is widely used in convolutional neural networks. This operation
+    applies Batch Normalization over inputs to avoid internal covariate shift as described
+    in the paper `Batch Normalization: Accelerating Deep Network Training by Reducing Internal
+    Covariate Shift <https://arxiv.org/abs/1502.03167>`_. It rescales and recenters the
+    features using a mini-batch of data and the learned parameters can be described
+    in the following formula,
+
+    .. math::
+
+        y = \frac{x - mean}{\sqrt{variance + \epsilon}} * \gamma + \beta
+
+    where :math:`\gamma` is `weight`, :math:`\beta` is `bias`, :math:`\epsilon` is `eps`, :math:`mean` is the
+    mean of :math:`x`, :math:`variance` is the variance of :math:`x`.
+
+    .. warning::
+        - For Atlas 200/300/500 inference product,
+          the result accuracy fails to reach 1‰ due to the square root instruction.
+
+    Note:
+        - If `training` is `False`, `weight`, `bias`, `running_mean` and `running_var` are Tensors.
+        - If `training` is `True`, `weight`, `bias`, `running_mean` and `running_var` are Parameters.
+
+    Args:
+        input (Tensor): Tensor of shape :math:`(N, C, *)`, with bfloat16, float16 or float32 data type.
+        running_mean (Union[Tensor, Parameter]): The shape :math:`(C,)`, has the same data type with `weight`.
+        running_var (Union[Tensor, Parameter]): The shape :math:`(C,)`, has the same data type with `weight`.
+        weight (Union[Tensor, Parameter]): The shape :math:`(C,)`, with bfloat, float16 or float32 data type.
+        bias (Union[Tensor, Parameter]): The shape :math:`(C,)`, has the same data type with `weight`.
+        training (bool, optional): If `training` is `True`, `mean` and `variance` are computed during training.
+            If `training` is `False`, they're loaded from checkpoint during inference. Default: ``False`` .
+        momentum (float, optional): The hyper parameter to compute moving average for `running_mean` and `running_var`
+            (e.g. :math:`new\_running\_mean = (1 - momentum) * running\_mean + momentum * current\_mean`).
+            Default: ``0.1`` .
+        eps (float, optional): A small value added for numerical stability. Default: ``1e-5``.
+
+    Returns:
+        output_x (Tensor) - The same type and shape as the `input_x`. The shape is :math:`(N, C, *)`.
+
+    Raises:
+        TypeError: If `training` is not a bool.
+        TypeError: If dtype of `eps` or `momentum` is not float.
+        TypeError: If `input_x`, `weight`, `bias`, `running_mean` or `running_var` is not a Tensor.
+        TypeError: If dtype of `input_x`, `weight` is not bfloat16, float16 or float32.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, ops
+        >>> input_x = Tensor([[1.0, 2.0], [3.0, 4.0]], mindspore.float32)
+        >>> running_mean = Tensor([0.5, 1.5], mindspore.float32)
+        >>> running_var = Tensor([0.1, 0.2], mindspore.float32)
+        >>> weight = Tensor([2.0, 2.0], mindspore.float32)
+        >>> bias = Tensor([-1.0, -1.0], mindspore.float32)
+        >>> output = ops.batch_norm(input_x, running_mean, running_var, weight, bias)
+        >>> print(output)
+        [[ 2.1621194  1.2360122]
+         [14.810596  10.180061 ]]
+    """
+    if weight is None:
+        weight = ops.ones([input.shape[1]], dtype=input.dtype)
+    if bias is None:
+        bias = ops.zeros([input.shape[1]], dtype=input.dtype)
+    batch_norm_op = BatchNormExt(training=training, momentum=momentum, eps=eps)
+    output = batch_norm_op(input, weight, bias, running_mean, running_var)
+    return output[0]
+
+
+
+__all__ = ['conv2d', 'max_pool2d', 'leaky_relu_ext', 'batch_norm']

mindspore/ops/function/__init__.py

@@ -26,8 +26,7 @@ from . import (
     nn_func,
     linalg_func,
     clip_func,
-    fft_func,
-    reshard_func
+    fft_func
 )
 from .array_func import (
     unique,
@@ -100,6 +99,7 @@ from .array_func import (
     tensor_scatter_elements,
     scatter,
     scatter_add,
+    scatter_add_ext,
     scatter_mul,
     scatter_max,
     scatter_min,
@@ -138,7 +138,6 @@ from .array_func import (
     index_fill,
     index_select,
     max,
-    argmax,
     min,
     population_count,
     topk,
@@ -166,6 +165,9 @@ from .array_func import (
     top_k,
     deepcopy,
     arange_ext,
+    zeros_like_ext,
+    ones_like_ext,
+    full_ext,
 )
 from .parameter_func import (
     assign,
@@ -178,6 +180,7 @@ from .math_func import (
     addn,
     absolute,
     abs,
+    argmax,
     argmin,
     angle,
     bincount,
@@ -197,6 +200,7 @@ from .math_func import (
     le,
     lerp,
     norm,
+    norm_ext,
     vector_norm,
     matrix_norm,
     round,
@@ -264,6 +268,7 @@ from .math_func import (
     matrix_determinant,
     det,
     linspace,
+    linspace_ext,
     lu_solve,
     matrix_solve,
     maximum,
@@ -369,6 +374,7 @@ from .math_func import (
     amin,
     amax,
     mean,
+    mean_ext,
     prod,
     all,
     any,
@@ -456,7 +462,6 @@ from .nn_func import (
     max_pool2d,
     max_pool3d,
     batch_norm,
-    rms_norm,
     bidense,
     celu,
     bias_add,
@@ -505,6 +510,7 @@ from .nn_func import (
     softplus,
     pdist,
     pad,
+    pad_ext,
     prelu,
     mirror_pad,
     nll_loss,
@@ -543,6 +549,7 @@ from .nn_func import (
     channel_shuffle,
     hardsigmoid,
     group_norm,
+    dropout_ext,
 )
 from .linalg_func import (
     cond,
@@ -593,6 +600,7 @@ from .random_func import (
     standard_laplace,
     random_categorical,
     uniform,
+    uniform_ext,
     standard_normal,
     random_gamma,
     uniform_candidate_sampler,
@@ -600,6 +608,7 @@ from .random_func import (
     log_uniform_candidate_sampler,
     shuffle,
     choice_with_mask,
+    normal_ext,
     normal,
     laplace,
     gamma,
@@ -733,9 +742,6 @@ from .other_func import (
     depend,
     partial,
 )
-from .reshard_func import (
-    reshard,
-)
 
 from ..operations.manually_defined import (rank, scalar_cast)
 
@@ -756,5 +762,4 @@ __all__.extend(sparse_unary_func.__all__)
 __all__.extend(clip_func.__all__)
 __all__.extend(fft_func.__all__)
 __all__.extend(other_func.__all__)
-__all__.extend(reshard_func.__all__)
 __all__.sort()