mindspore 1.10.0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/ops/composite/env_ops.py

@@ -0,0 +1,41 @@
+# This is the Python adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
+#
+# Copyright 2022 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.
+# ============================================================================
+
+"""Env related operations."""
+from __future__ import absolute_import
+from mindspore.ops.composite.base import MultitypeFuncGraph
+from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
+from mindspore.ops.primitive import Primitive
+from mindspore.ops.operations import _grad_ops
+from mindspore.ops import operations as P
+
+env_get = MultitypeFuncGraph("env_get")
+environ_get = Primitive('EnvironGet')
+ref_to_embed = _grad_ops.RefToEmbed()
+tensor_zeros_like = P.ZerosLike()
+
+
+@env_get.register("EnvType", "Tensor")
+def _tensor_env_get(env, parameter):
+    """Used to get env."""
+    return environ_get(env, ref_to_embed(parameter), tensor_zeros_like(parameter))
+
+
+@env_get.register("EnvType", "MapTensor")
+def _map_tensor_env_get(env, map_parameter):
+    """Used to get env for map parameter."""
+    return environ_get(env, ref_to_embed(map_parameter), zeros_like(map_parameter))

mindspore/ops/composite/math_ops.py

@@ -13,22 +13,22 @@
 # limitations under the License.
 # ============================================================================
 """math Operations."""
-import numpy as np
 from mindspore.ops.composite.multitype_ops import _constexpr_utils as const_utils
 from mindspore.common import dtype as mstype
-from mindspore._checkparam import Validator as validator
-from mindspore.ops.primitive import constexpr
+from mindspore import _checkparam as validator
+from mindspore.ops.primitive import constexpr, _primexpr
 from mindspore.ops import functional as F
-from mindspore.ops.operations._inner_ops import DynamicResizeNearestNeighbor
-from ..function.math_func import cummin as cummin_
-from .. import operations as P
+from mindspore.ops.function.math_func import cummin as cummin_
+from mindspore.ops import operations as P
 
 
-@constexpr
+@_primexpr
 def _check_validate_axis(axis, name):
-    if isinstance(axis, (tuple, list)):
-        for idx, item in enumerate(axis):
-            validator.check_value_type("axis[%d]" % idx, item, [int], name)
+    def _check(axis):
+        if isinstance(axis, (tuple, list)):
+            for idx, item in enumerate(axis):
+                validator.check_value_type("axis[%d]" % idx, item, [int], name)
+    _check(axis)
     axis = validator.check_value_type('axis', axis, [int, tuple, list], name)
     return axis
 
@@ -46,20 +46,26 @@ def is_const(x):
 
 def count_nonzero(x, axis=(), keep_dims=False, dtype=mstype.int32):
     r"""
-    Count number of nonzero elements across axis of input tensor
+    Count number of nonzero elements across axis of input tensor.
 
     Args:
-        x (Tensor): Input data is used to count non-zero numbers.
-          :math:`(N,*)` where :math:`*` means, any number of additional dimensions.
-        axis (Union[int, tuple(int), list(int)]): The dimensions to reduce. Only constant value is allowed.
-                                                  Default: (), reduce all dimensions.
-        keep_dims (bool): If true, keep these reduced dimensions and the length is 1.
-                          If false, don't keep these dimensions. Default: False.
-        dtype (Union[Number, mindspore.bool\_]): The data type of the output tensor. Only constant value is allowed.
-                                             Default: mindspore.int32
+        x (Tensor): Input data is used to count non-zero numbers. With shape
+            :math:`(N,*)` where :math:`*` means, any number of additional dimensions.
+        axis (Union[int, tuple(int), list(int)], optional): The dimensions to reduce.
+            Default: (), reduce all dimensions.
+        keep_dims (bool, optional): Whether to maintain dimensions specified by `axis`.
+            If true, keep these reduced dimensions and the length is 1.
+            If false, don't keep these dimensions. Default: False.
+        dtype (Union[Number, mindspore.bool\_], optional): The data type of the output tensor.
+            Default: mindspore.int32.
 
     Returns:
-          Tensor, number of nonzero element. The data type is `dtype`.
+          Tensor, number of nonzero element across axis specified by `axis`.
+          The data type is specified by `dtype`.
+
+    Raises:
+        TypeError: If `axis` is not int, tuple or list.
+        ValueError: If any value in `axis` is not in range [-x.ndim, x.ndim).
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -102,7 +108,9 @@ def count_nonzero(x, axis=(), keep_dims=False, dtype=mstype.int32):
     not_equal = P.NotEqual()
     cast = P.Cast()
     reduce_sum = P.ReduceSum(keep_dims)
-    nonzero_bool = not_equal(x, 0)
+    zeros = P.Zeros()
+    tensor_0 = zeros(x.shape, x.dtype)
+    nonzero_bool = not_equal(x, tensor_0)
     # ReduceSum only support float16 or float32 tensor.
     nonzero_val = cast(nonzero_bool, mstype.float32)
     nonzero_num = cast(reduce_sum(nonzero_val, axis), dtype)
@@ -110,7 +118,7 @@ def count_nonzero(x, axis=(), keep_dims=False, dtype=mstype.int32):
     return nonzero_num
 
 
-@constexpr
+@_primexpr
 def _int_to_tuple_conv(axes):
     """
     Converts ints to tuples in input axes, expected by most validation checks.
@@ -121,7 +129,7 @@ def _int_to_tuple_conv(axes):
     return axes
 
 
-@constexpr
+@_primexpr
 def _check_axes(axes, prim_name=None):
     """
     Check for validity and type of axes passed to function.
@@ -154,21 +162,29 @@ def _typecheck_input(x1_type, x2_type, prim_name=None):
                         f"and x2_type: {x2_type}.")
 
 
-@constexpr
+@_primexpr
 def _axes_int_check(x1_shape, x2_shape, axes, prim_name=None):
     """
     Convert from single int axes to 2d tuple if required
     """
     msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
-    if isinstance(axes, int):
+
+    def _check_lt_zero(axes):
         if axes < 0:
             raise ValueError(f"{msg_prefix} 'axes' must be at least 0, but got {axes}.")
-        if axes == 0:
-            # outer product, no input validation required
-            return [], []
+
+    def _check_len(axes, x1_shape, x2_shape):
         if axes > len(x1_shape) or axes > len(x2_shape):
             raise ValueError(f"{msg_prefix} 'axes' cannot be greater than the length of 'x1_shape' and 'x2_shape', "
                              f"but got 'axes': {axes}, 'x1_shape': {x1_shape}, 'x2_shape': {x2_shape}.")
+
+
+    if isinstance(axes, int):
+        _check_lt_zero(axes)
+        if axes == 0:
+            # outer product, no input validation required
+            return [], []
+        _check_len(axes, x1_shape, x2_shape)
         x1_ind = tuple(range(len(x1_shape))[-1 * axes:])
         x2_ind = tuple(range(len(x2_shape))[:axes])
         axes = tuple((x1_ind, x2_ind))
@@ -176,7 +192,7 @@ def _axes_int_check(x1_shape, x2_shape, axes, prim_name=None):
     return axes
 
 
-@constexpr
+@_primexpr
 def _validate_axes(x1_shape, x2_shape, axes, prim_name=None):
     """
     Checks for axes having the correct length according to input, for any value in axis
@@ -184,25 +200,32 @@ def _validate_axes(x1_shape, x2_shape, axes, prim_name=None):
     with given inputs.
     """
     msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
+
+    def _check_len(axes_len, shape_dim_len, x_axes):
+        if axes_len > shape_dim_len:
+            raise ValueError(f"{msg_prefix} length of element {x_axes} in 'axes' must be less than or equal to "
+                             f"{shape_dim_len}, but got {axes_len}.")
+
+    def _check_value(x_axes, min_val, max_val):
+        for _, x_value in enumerate(x_axes):
+            if x_value > max_val or x_value < min_val:
+                raise ValueError(f"{msg_prefix} value in 'axes' must be in range: [{min_val}, {max_val}], "
+                                 f"but got {x_value}.")
+
     shapes = [x1_shape, x2_shape]
 
     # axis length check
     for ix_input, x_axes in enumerate(axes):
         axes_len = len(x_axes)
         shape_dim_len = len(shapes[ix_input])
-        if axes_len > shape_dim_len:
-            raise ValueError(f"{msg_prefix} length of element {x_axes} in 'axes' must be less than or equal to "
-                             f"{shape_dim_len}, but got {axes_len}.")
+        _check_len(axes_len, shape_dim_len, x_axes)
 
     # axis values range check
     for ix_input, x_axes in enumerate(axes):
         comp_shape = shapes[ix_input]
         max_val = len(comp_shape) - 1
         min_val = -1 * len(comp_shape)
-        for _, x_value in enumerate(x_axes):
-            if not min_val <= x_value <= max_val:
-                raise ValueError(f"{msg_prefix} value in 'axes' must be in range: [{min_val}, {max_val}], "
-                                 f"but got {x_value}.")
+        _check_value(x_axes, min_val, max_val)
 
     # check axis value with input shape - both ways for axis valid
     invalid_a = False
@@ -212,23 +235,31 @@ def _validate_axes(x1_shape, x2_shape, axes, prim_name=None):
             invalid_a = True
         if x1_shape[axes[0][i]] != x2_shape[axes[1][len(axes[0]) - 1 - i]]:
             invalid_b = True
-    if invalid_a and invalid_b:
-        raise ValueError(f"{msg_prefix} 'i' should exist such that 'x1_shape[axes[0][i]]' is equal to "
-                         f"'x2_shape[axes[1][i]]' or 'x2_shape[axes[1][len(axes[0])-1-i]]', but got "
-                         f"'x1_shape': {x1_shape}, 'x2_shape': {x2_shape}, 'axes': {axes}.")
 
+    def _check(invalid_a, invalid_b, x1_shape, x2_shape, axes):
+        if invalid_a and invalid_b:
+            raise ValueError(f"{msg_prefix} 'i' should exist such that 'x1_shape[axes[0][i]]' is equal to "
+                             f"'x2_shape[axes[1][i]]' or 'x2_shape[axes[1][len(axes[0])-1-i]]', but got "
+                             f"'x1_shape': {x1_shape}, 'x2_shape': {x2_shape}, 'axes': {axes}.")
 
-@constexpr
+    _check(invalid_a, invalid_b, x1_shape, x2_shape, axes)
+
+
+@_primexpr
 def _calc_new_shape(shape, axes, position=0):
     """
     Calculate transpose and reshape parameters for input transformations,
     'position' refers to whether tensor is first or second in the op.
     """
     contraction_axes = tuple(i if i >= 0 else i + len(shape) for i in axes[position])
-    prod_contraction = int(np.prod([shape[i] for i in contraction_axes]))
+    prod_contraction = 1
+    for i in contraction_axes:
+        prod_contraction *= shape[i]
     free_axes = tuple(i for i in range(len(shape)) if i not in contraction_axes)
-    free_dims = tuple(shape[i] for i in free_axes)
-    prod_free = int(np.prod(free_dims))
+    free_dims = tuple(shape[i] if shape[i] is not None else -1 for i in free_axes)
+    prod_free = 1
+    for free_dim in free_dims:
+        prod_free *= free_dim
 
     transpose_perm = contraction_axes + free_axes if position else free_axes + contraction_axes
     new_shape = (prod_contraction, prod_free) if position else (prod_free, prod_contraction)
@@ -245,7 +276,7 @@ def tensor_dot(x1, x2, axes):
 
     Selected dims in both inputs must also match.
 
-    axes = 0 leads to outer product
+    axes = 0 leads to outer product.
     axes = 1 leads to normal matrix multiplication when inputs both 2D.
     axes = 1 is the same as axes = ((1,),(0,)) where both `a` and `b` are 2D.
     axes = 2 is the same as axes = ((1,2),(0,1)) where both `a` and `b` are 3D.
@@ -288,10 +319,7 @@ def tensor_dot(x1, x2, axes):
     # input validity checks
     x1_shape = shape_op(x1)
     x2_shape = shape_op(x2)
-    x1_type = F.dtype(x1)
-    x2_type = F.dtype(x2)
     axes = _check_axes(axes, 'tensor_dot')
-    _typecheck_input(x1_type, x2_type, 'tensor_dot')
     # input compatibility check & axes format update
     axes = _axes_int_check(x1_shape, x2_shape, axes, 'tensor_dot')
     _validate_axes(x1_shape, x2_shape, axes, 'tensor_dot')
@@ -308,7 +336,7 @@ def tensor_dot(x1, x2, axes):
     return final_result
 
 
-@constexpr
+@_primexpr
 def _check_invalid_input(x1_shape, x2_shape, prim_name=None):
     msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
     if len(x1_shape) < 2 or len(x2_shape) < 2:
@@ -329,56 +357,57 @@ def _typecheck_input_dot(x1_type, x2_type, prim_name=None):
                         f"x1_type: {x1_type} and x2_type: {x2_type}.")
 
 
-@constexpr
+@_primexpr
 def _get_transpose_shape(x2_shape):
     x2_shape_range = tuple(range(len(x2_shape)))
     x2_shape_transpose = x2_shape_range[-2:-1] + x2_shape_range[:-2] + x2_shape_range[-1:]
     return x2_shape_transpose
 
 
-def dot(x1, x2):
+def dot(input, other):
     """
     Computation a dot product between samples in two tensors.
 
     Args:
-        x1 (Tensor): First tensor in Dot op with datatype float16 or float32,
+        input (Tensor): First tensor in Dot op with datatype float16 or float32,
             The rank must be greater than or equal to 2.
-        x2 (Tensor): Second tensor in Dot op with datatype float16 or float32,
+        other (Tensor): Second tensor in Dot op with datatype float16 or float32,
             The rank must be greater than or equal to 2.
 
     Returns:
-        Tensor, dot product of x1 and x2.
+        Tensor, dot product of input and other.
 
     Raises:
-        TypeError: If type of x1 and x2 are not the same.
-        TypeError: If dtype of x1 or x2 is not float16 or float32.
-        ValueError: If rank of x1 or x2 less than 2.
+        TypeError: If type of input and other are not the same.
+        TypeError: If dtype of input or other is not float16 or float32.
+        ValueError: If rank of input or other less than 2.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
+        >>> import numpy as np
         >>> import mindspore
         >>> from mindspore import Tensor, ops
-        >>> input_x1 = Tensor(np.ones(shape=[2, 3]), mindspore.float32)
-        >>> input_x2 = Tensor(np.ones(shape=[1, 3, 2]), mindspore.float32)
-        >>> output = ops.dot(input_x1, input_x2)
+        >>> input = Tensor(np.ones(shape=[2, 3]), mindspore.float32)
+        >>> other = Tensor(np.ones(shape=[1, 3, 2]), mindspore.float32)
+        >>> output = ops.dot(input, other)
         >>> print(output)
         [[[3. 3.]]
          [[3. 3.]]]
         >>> print(output.shape)
         (2, 1, 2)
-        >>> input_x1 = Tensor(np.ones(shape=[1, 2, 3]), mindspore.float32)
-        >>> input_x2 = Tensor(np.ones(shape=[1, 3, 2]), mindspore.float32)
-        >>> output = ops.dot(input_x1, input_x2)
+        >>> input = Tensor(np.ones(shape=[1, 2, 3]), mindspore.float32)
+        >>> other = Tensor(np.ones(shape=[1, 3, 2]), mindspore.float32)
+        >>> output = ops.dot(input, other)
         >>> print(output)
         [[[[3. 3.]]
           [[3. 3.]]]]
         >>> print(output.shape)
         (1, 2, 1, 2)
-        >>> input_x1 = Tensor(np.ones(shape=[1, 2, 3]), mindspore.float32)
-        >>> input_x2 = Tensor(np.ones(shape=[2, 3, 2]), mindspore.float32)
-        >>> output = ops.dot(input_x1, input_x2)
+        >>> input = Tensor(np.ones(shape=[1, 2, 3]), mindspore.float32)
+        >>> other = Tensor(np.ones(shape=[2, 3, 2]), mindspore.float32)
+        >>> output = ops.dot(input, other)
         >>> print(output)
         [[[[3. 3.]
            [3. 3.]]
@@ -386,9 +415,9 @@ def dot(x1, x2):
            [3. 3.]]]]
         >>> print(output.shape)
         (1, 2, 2, 2)
-        >>> input_x1 = Tensor(np.ones(shape=[3, 2, 3]), mindspore.float32)
-        >>> input_x2 = Tensor(np.ones(shape=[2, 1, 3, 2]), mindspore.float32)
-        >>> output = ops.dot(input_x1, input_x2)
+        >>> input = Tensor(np.ones(shape=[3, 2, 3]), mindspore.float32)
+        >>> other = Tensor(np.ones(shape=[2, 1, 3, 2]), mindspore.float32)
+        >>> output = ops.dot(input, other)
         >>> print(output)
         [[[[[3. 3.]]
            [[3. 3.]]]
@@ -409,34 +438,36 @@ def dot(x1, x2):
     reshape_op = P.Reshape()
     transpose_op = P.Transpose()
     matmul_op = P.MatMul(False, False)
-    x1_shape = shape_op(x1)
-    x2_shape = shape_op(x2)
-    x1_type = F.dtype(x1)
-    x2_type = F.dtype(x2)
-    _typecheck_input_dot(x1_type, x2_type, 'dot')
-    _check_invalid_input(x1_shape, x2_shape, 'dot')
-
-    if len(x1_shape) > 2 or len(x2_shape) > 2:
-        x2_shape_transpose = _get_transpose_shape(x2_shape)
-        x2_transpose = transpose_op(x2, x2_shape_transpose)
-        x1_reshape = reshape_op(x1, (-1, x1_shape[-1]))
-        x2_reshape = reshape_op(x2_transpose, (x2_shape[-2], -1))
-        mul_result = matmul_op(x1_reshape, x2_reshape)
-        reshape_shape = x1_shape[:-1] + x2_shape[:-2] + x2_shape[-1:]
+    input_shape = shape_op(input)
+    other_shape = shape_op(other)
+    input_type = F.dtype(input)
+    other_type = F.dtype(other)
+    _typecheck_input_dot(input_type, other_type, 'dot')
+    _check_invalid_input(input_shape, other_shape, 'dot')
+
+    if len(input_shape) > 2 or len(other_shape) > 2:
+        other_shape_transpose = _get_transpose_shape(other_shape)
+        other_transpose = transpose_op(other, other_shape_transpose)
+        input_reshape = reshape_op(input, (-1, input_shape[-1]))
+        other_reshape = reshape_op(other_transpose, (other_shape[-2], -1))
+        mul_result = matmul_op(input_reshape, other_reshape)
+        reshape_shape = input_shape[:-1] + other_shape[:-2] + other_shape[-1:]
         reshape_shape = (-1,) + reshape_shape[1:]
         return reshape_op(mul_result, reshape_shape)
-    return matmul_op(x1, x2)
+    return matmul_op(input, other)
 
 
-@constexpr
+@_primexpr
 def _get_batch_size(x1_shape, x2_shape, prim_name=None):
     """
     Get batch sizes from two inputs
     """
-    msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
-    if len(x1_shape) < 2 or len(x2_shape) < 2:
-        raise ValueError(f"{msg_prefix} inputs x1, x2 should have 'dimension >= 2', "
-                         f"but got 'len(x1_shape)': ({len(x1_shape)}) and 'len(x2_shape)': ({len(x2_shape)}).")
+    def _check():
+        msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
+        if len(x1_shape) < 2 or len(x2_shape) < 2:
+            raise ValueError(f"{msg_prefix} inputs x1, x2 should have 'dimension >= 2', "
+                             f"but got 'len(x1_shape)': ({len(x1_shape)}) and 'len(x2_shape)': ({len(x2_shape)}).")
+    _check()
     return x1_shape[0], x2_shape[0]
 
 
@@ -453,12 +484,33 @@ def _typecheck_input_batch_dot(x1_type, x2_type, prim_name=None):
                         f"x2_type: {x2_type}.")
 
 
-@constexpr
+@_primexpr
 def _check_axes_for_batch_dot(x1_shape, x2_shape, axes, prim_name=None):
     """
     Check whether axes are valid and cast axes from tuple to list
     """
     msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
+
+    def _check_1(axes):
+        if 0 in axes:
+            raise ValueError(f"{msg_prefix} 'axes' cannot contain 0, but got axes: {axes}.")
+        if len(axes) != 2:
+            raise ValueError(f"{msg_prefix} length of 'axes' must be equal to 2, but got {len(axes)}.")
+
+    def _check_2(axes, x1_shape, x2_shape):
+        if axes[0] > len(x1_shape) or axes[1] > len(x2_shape):
+            raise ValueError(f"{msg_prefix} axes[0] must be less than or equal to len(x1_shape), "
+                             f"and axes[1] must be less than or equal to len(x2_shape)."
+                             f"But got 'axes': {axes}, 'x1_shape': {x1_shape}, 'x2_shape': {x2_shape}.")
+
+    def _check_3(axes, x1_shape, x2_shape):
+        if axes == 0:
+            raise ValueError(f"{msg_prefix} 'axes' should not be equal to 0, but got {axes}.")
+
+        if axes > len(x1_shape) or axes > len(x2_shape):
+            raise ValueError(f"{msg_prefix} 'axes' cannot be greater than the length of 'x1_shape' and 'x2_shape', "
+                             f"but got 'axes': {axes}, 'x1_shape': {x1_shape}, 'x2_shape': {x2_shape}.")
+
     if axes is None:
         if len(x2_shape) == 2:
             axes = [len(x1_shape) - 1, len(x2_shape) - 1]
@@ -466,10 +518,7 @@ def _check_axes_for_batch_dot(x1_shape, x2_shape, axes, prim_name=None):
             axes = [len(x1_shape) - 1, len(x2_shape) - 2]
 
     if isinstance(axes, (list, tuple)):
-        if 0 in axes:
-            raise ValueError(f"{msg_prefix} 'axes' cannot contain 0, but got axes: {axes}.")
-        if len(axes) != 2:
-            raise ValueError(f"{msg_prefix} length of 'axes' must be equal to 2, but got {len(axes)}.")
+        _check_1(axes)
         if isinstance(axes, tuple):
             axes = list(axes)
         validator.check_value_type('axes[0]', axes[0], [int], 'batch_dot')
@@ -481,19 +530,12 @@ def _check_axes_for_batch_dot(x1_shape, x2_shape, axes, prim_name=None):
             axes[1] += len(x2_shape)
         validator.check_non_negative_int(axes[0], 'reversed axes[0]', 'batch_dot')
         validator.check_non_negative_int(axes[1], 'reversed axes[1]', 'batch_dot')
-        if axes[0] > len(x1_shape) or axes[1] > len(x2_shape):
-            raise ValueError(f"{msg_prefix} axes[0] must be less than or equal to len(x1_shape), "
-                             f"and axes[1] must be less than or equal to len(x2_shape)."
-                             f"But got 'axes': {axes}, 'x1_shape': {x1_shape}, 'x2_shape': {x2_shape}.")
+        _check_2(axes, x1_shape, x2_shape)
     elif isinstance(axes, int):
-        if axes == 0:
-            raise ValueError(f"{msg_prefix} 'axes' should not be equal to 0, but got {axes}.")
+        _check_3(axes, x1_shape, x2_shape)
         if axes < 0:
             axes = [axes + len(x1_shape), axes + len(x2_shape)]
             validator.check_non_negative_int(axes[0], 'reversed axes', 'batch_dot')
-        elif axes > len(x1_shape) or axes > len(x2_shape):
-            raise ValueError(f"{msg_prefix} 'axes' cannot be greater than the length of 'x1_shape' and 'x2_shape', "
-                             f"but got 'axes': {axes}, 'x1_shape': {x1_shape}, 'x2_shape': {x2_shape}.")
         else:
             axes = [axes, axes]
     else:
@@ -502,7 +544,7 @@ def _check_axes_for_batch_dot(x1_shape, x2_shape, axes, prim_name=None):
     return axes
 
 
-@constexpr
+@_primexpr
 def _calc_new_shape_batchdot(shape, axes, position=0):
     """
     Calculate transpose and reshape parameters for input transformations,
@@ -510,10 +552,14 @@ def _calc_new_shape_batchdot(shape, axes, position=0):
     """
     axis = axes[position]
     contraction_axes = tuple([axis])
-    prod_contraction = int(np.prod([shape[i] for i in contraction_axes]))
+    prod_contraction = 1
+    for i in contraction_axes:
+        prod_contraction *= shape[i]
     free_axes = tuple(i for i in range(1, len(shape)) if i not in contraction_axes)
     free_dims = tuple(shape[i] for i in free_axes)
-    prod_free = int(np.prod(free_dims))
+    prod_free = 1
+    for free_dim in free_dims:
+        prod_free *= free_dim
 
     transpose_perm = contraction_axes + free_axes if position else free_axes + contraction_axes
     transpose_perm = tuple([0]) + transpose_perm
@@ -522,7 +568,7 @@ def _calc_new_shape_batchdot(shape, axes, position=0):
     return new_shape, transpose_perm, free_dims
 
 
-@constexpr
+@_primexpr
 def _check_batch_size(x1_batch_size, x2_batch_size, prim_name=None):
     """
     Check whether batch size of two inputs are the same
@@ -533,7 +579,7 @@ def _check_batch_size(x1_batch_size, x2_batch_size, prim_name=None):
                          f"'x1_batch_size': {x1_batch_size} and 'x2_batch_size': {x2_batch_size}.")
 
 
-@constexpr
+@_primexpr
 def _get_output_shape(batch_size, x1_ret, x2_ret):
     """
     Compute output shape for batch dot
@@ -725,6 +771,49 @@ def matmul(x1, x2, dtype=None):
     return res
 
 
+def mm(input, mat2):
+    r"""
+    Returns the matrix product of two arrays.
+    If `input` is a :math:`(n \times m)` Tensor, `mat2` is a
+    :math:`(m \times p)` Tensor, `out` will be a :math:`(n \times p)` Tensor.
+
+    Note:
+        This function cannot support broadcasting.
+        Refer to :func:`mindspore.ops.matmul` instead if you need a broadcastable function.
+
+    Args:
+        input (Tensor): The first matrix of matrix multiplication.
+          The last dimension of `input` must be the same size as the first dimension of `mat2`.
+        mat2 (Tensor): The second matrix of matrix multiplication.
+          The last dimension of `input` must be the same size as the first dimension of `mat2`.
+
+    Returns:
+        Tensor or scalar, the matrix product of the inputs.
+
+    Raises:
+        ValueError: If the last dimension of `input` is not the same size as the
+            second-to-last dimension of `mat2`.
+        ValueError: If `input` or `mat2` is not a matrix.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> import mindspore.ops as ops
+        >>> import numpy as np
+        >>> x1 = ms.Tensor(np.random.rand(2, 3))
+        >>> x2 = ms.Tensor(np.random.rand(3, 4))
+        >>> out = ops.mm(x1, x2)
+        >>> print(out.shape)
+        (2, 4)
+    """
+    if input.ndim != 2 or mat2.ndim != 2:
+        raise ValueError(f"For mm, the input tensor must be a matrix, "
+                         f"but got mat1.ndim:{input.ndim}, mat2.ndim:{mat2.ndim}")
+    return matmul(input, mat2)
+
+
 def cummin(x, axis):
     r"""
     Returns a tuple (values,indices) where 'values' is the cumulative minimum value of input Tensor `x`
@@ -763,51 +852,3 @@ def cummin(x, axis):
         [0 1 1 1 4 4]
     """
     return cummin_(x, axis)
-
-
-def resize_nearest_neighbor(input_x, size, align_corners=False):
-    r"""
-    Resizes the input tensor by using the nearest neighbor algorithm.
-
-    Resizes the input tensor to a given size by using the nearest neighbor algorithm. The nearest
-    neighbor algorithm selects the value of the nearest point and does not consider the
-    values of neighboring points at all, yielding a piecewise-constant interpolant.
-
-    Args:
-        input_x (Tensor) - The input tensor. The shape of the tensor is :math:`(N, C, H, W)`.
-        size (Union[Tensor, tuple, list]): The target size. The dimension of size must be 2.
-        align_corners (bool): Whether the centers of the 4 corner pixels of the input
-                              and output tensors are aligned. Default: False.
-
-    Returns:
-        Tensor, the shape of the output tensor is  :math:`(N, C, NEW\_H, NEW\_W)`.
-        The data type is the same as the `input_x`.
-
-    Raises:
-        TypeError: If `input_x` is not a Tensor.
-        TypeError: If `size` is neither tuple nor list.
-        TypeError: If `align_corners` is not a bool.
-        ValueError: If length of `size` is not equal to 2.
-
-    Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
-
-    Examples:
-        >>> input_tensor = Tensor(np.array([[[[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]]]]), mindspore.float32)
-        >>> output = ops.ResizeNearestNeighbor(input_tensor, (2, 2))
-        >>> print(output)
-        [[[[-0.1  0.3]
-           [ 0.4  0.5]]]]
-    """
-    if size is None:
-        raise ValueError(f'For ResizeNearestNeighbor, size could not be None.')
-    if isinstance(size, (tuple, list)):
-        resize = P.ResizeNearestNeighbor(size, align_corners)
-        return resize(input_x)
-    if is_const(size):
-        size = size.asnumpy()
-        resize = P.ResizeNearestNeighbor(size, align_corners)
-        return resize(input_x)
-
-    resize = DynamicResizeNearestNeighbor(align_corners)
-    return resize(input_x, size)