mindspore 2.1.0__cp38-cp38-manylinux1_x86_64.whl → 2.2.11__cp38-cp38-manylinux1_x86_64.whl

mindspore/ops/_op_impl/tbe/trans_data_ds.py

@@ -28,6 +28,8 @@ trans_data_op_info = TBERegOp("TransData") \
           "DefaultFormat, NC1HWC0, FRACTAL_Z, FRACTAL_NZ, HWCN, C1HWNCoC0, NDHWC, NHWC") \
     .attr("dst_format", "required", "str",
           "DefaultFormat, NC1HWC0, FRACTAL_Z, FRACTAL_NZ, HWCN, C1HWNCoC0, NDHWC, NHWC") \
+    .attr("src_subformat", "optional", "int", "all", "1") \
+    .attr("dst_subformat", "optional", "int", "all", "1") \
     .attr("groups", "optional", "int", "all", "1") \
     .input(0, "src", False, "required", "all") \
     .output(0, "dst", False, "required", "all") \

mindspore/ops/_primitive_cache.py

@@ -85,6 +85,6 @@ def _get_cache_prim(cls: Primitive) -> Primitive:
             _PRIM_CACHE[key] = prim
         return _PRIM_CACHE.get(key)
 
-    if _is_need_compile(_temp_func):
+    if _is_need_compile(_temp_func): # @jit.cond: True
         return _new_prim_for_graph
     return _get_cache_prim_for_pynative

mindspore/ops/_tracefunc.py

@@ -17,12 +17,13 @@ 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
+from mindspore.common.api import _handle_func_args, _pynative_executor
 
 
 class _PackTensor(StubTensor):
@@ -64,6 +65,7 @@ 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__)
@@ -79,19 +81,29 @@ class PackFunc(Primitive):
                 args = (self.cell_obj, *args)
             return self.func(*args, **kwargs)
         self.kwargs = kwargs
-        return super().__call__(*args)
+        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)
-            return ret
-        return self._run_op(args)
+        else:
+            ret = self._run_op(args)
+        PackFunc.current = old
+        return ret
 
     @staticmethod
     def is_tracing():
-        return _RunOpHook.current and _RunOpHook.current.hook is PackFunc._trace_run_op
+        return PackFunc.current is not None
 
     @staticmethod
     def _trace_run_op(obj, args):
@@ -197,13 +209,33 @@ def trace(fn):
 
     @functools.wraps(fn)
     def _trace_wrap(*args, **kwargs):
-        args, kwargs = _handle_func_args(fn, *args, **kwargs)
-        obj = None
-
-        if args and not isinstance(args[0], Tensor) and hasattr(args[0], fn.__name__):
-            obj, args = args[0], args[1:]
-        key = f"{id(obj)}_{id(fn)}"
-
-        return PackFunc(fn, key, obj, True)(*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

@@ -78,6 +78,11 @@ def get_broadcast_shape(x_shape, y_shape, prim_name, arg_name1="x", arg_name2="y
     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)
@@ -98,7 +103,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)

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)

mindspore/ops/_vmap/vmap_base.py

@@ -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,

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_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)

mindspore/ops/_vmap/vmap_nn_ops.py

@@ -325,9 +325,10 @@ def get_bce_with_logits_loss_vamp_rule(prim, axis_size):
         # If rank is larger than 1, we need to reduce result when reduction != 'none'
         if max_rank > 1:
             reduce_indexes = tuple(range(1, max_rank))
-        if logits_dim == label_dim and F.shape(logits) == F.shape(label) \
-                and logits_dim == weight_dim and F.shape(logits) == F.shape(weight) \
-                and logits_dim == pos_weight_dim and F.shape(logits) == F.shape(pos_weight):
+        logits_dim_ok = logits_dim == label_dim and logits_dim == weight_dim and logits_dim == pos_weight_dim
+        shape = F.shape(logits)
+        shape_ok = shape == F.shape(label) and shape == F.shape(weight) and shape == F.shape(pos_weight)
+        if logits_dim_ok and shape_ok:
             if prim_reduction == 'none':
                 output = prim(logits, label, weight, pos_weight)
             elif prim_reduction in ('mean', 'sum'):
@@ -798,7 +799,8 @@ def get_instance_norm_rule(prim, axis_size):
             output_x, updated_moving_mean, updated_moving_variance = prim(input_x, gamma, beta, mean, variance, u_monad)
             return (output_x, None), (updated_moving_mean, None), (updated_moving_variance, None)
 
-        if gamma_dim != 0 or beta_dim != gamma_dim or mean_dim != gamma_dim or variance_dim != gamma_dim:
+        precondition = gamma_dim != 0 or beta_dim != gamma_dim or mean_dim != gamma_dim or variance_dim != gamma_dim
+        if precondition:
             # pylint: disable=too-many-format-args
             raise ValueError(
                 "For `{}`, the source axis of `var` must be equal to `accum` and `accum_update`, and not equal to 0, "
@@ -1679,7 +1681,8 @@ def get_rmsprop_vmap_rule(prim, axis_size):
             res = prim(var, mean_square, moment, lr, grad, decay, momentum, epsilon,
                        u_monad)  # low dimensional operator;
             return (res, None)
-        if var_dim != 0 or var_dim != mean_square_dim or var_dim != moment_dim or var_dim != grad_dim:
+        precondition = var_dim != 0 or var_dim != mean_square_dim or var_dim != moment_dim or var_dim != grad_dim
+        if precondition:
             raise ValueError(
                 f"For '{prim_name}', the source axis of 'var' must be equal to 'mean_square_dim' "
                 f"and 'moment_dim' and 'grad_dim' and not equal to 0, "
@@ -1735,8 +1738,8 @@ def get_apply_centered_rmsprop_vmap_rule(prim, axis_size):
             var = prim(var, mean_grad, mean_square,
                        mom, grad, lr, rho, momentum, eps, u_monad)
             return (var, None)
-
-        if var_dim != 0 or var_dim != mean_grad_dim or var_dim != mean_square_dim or var_dim != mom_dim:
+        precondition = var_dim != 0 or var_dim != mean_grad_dim or var_dim != mean_square_dim or var_dim != mom_dim
+        if precondition:
             raise ValueError(
                 f"For '{prim_name}', the source axis of 'var' must be equal to 'mean_grad_dim' "
                 f"and 'mean_square_dim' and 'mom_dim' and not equal to 0, "
@@ -2000,6 +2003,57 @@ def get_sparse_apply_ftrl_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register(P.Dense)
+def get_dense_vmap_rule(prim, axis_size):
+    """VmapRule for `Dense` operation."""
+    if isinstance(prim, str):
+        prim = Primitive(prim)
+
+    batch_matmul = P.BatchMatMul(transpose_b=True)
+
+    @_primexpr
+    def get_start_mid_end(x_shape):
+        start = x_shape[0]
+        mid = 1
+        for shp in x_shape[1:-1]:
+            mid *= shp
+        end = x_shape[-1]
+        return start, mid, end
+
+    def vmap_rule(x_bdim, w_bdim, b_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim, w_bdim, b_bdim)
+        if is_all_none:
+            return result
+
+        x, x_dim = x_bdim
+        w, w_dim = w_bdim
+        b, b_dim = b_bdim
+        x = _bdim_at_front(x, x_dim, axis_size)
+        w = _bdim_at_front(w, w_dim, axis_size)
+        if b is not None:
+            b = _bdim_at_front(b, b_dim, axis_size)
+
+        x_shape = x.shape
+        start, mid, end = get_start_mid_end(x_shape)
+
+        x = x.reshape(start, mid, end)
+
+        out = batch_matmul(x, w)
+        out_shape = tuple(x_shape[:-1]) + (out.shape[-1],)
+        out = out.reshape(out_shape)
+
+        if b is not None:
+            b_shape = b.shape
+            b_shape = (start,) + (1,) * (len(out_shape) - 2) + (b_shape[-1],)
+            b = b.reshape(b_shape)
+
+            out = out + b
+
+        return out, 0
+
+    return vmap_rule
+
+
 # Unary vmap
 get_unop_vmap_rule = vmap_rules_getters.register(P.Elu)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.ReLU)(get_unop_vmap_rule)

mindspore/ops/arg_dtype_cast.py

@@ -0,0 +1,54 @@
+# This is the Python adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
+#
+# Copyright 2023-2024 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.
+# ============================================================================
+"""Operator argument data type cast function."""
+from enum import Enum
+
+
+class TypeCastKind(Enum):
+    INT_TO_TUPLE = 1
+    INT_OR_TUPLE_TO_LIST = 2
+
+
+def type_it(src_data, cast_type):
+    """
+    cast operator argument data type.
+    """
+    if cast_type == TypeCastKind.INT_TO_TUPLE:
+        if isinstance(src_data, tuple):
+            return src_data
+
+        if isinstance(src_data, int):
+            return (src_data,)
+
+        raise TypeError(f'{src_data} is the wrong data type.')
+
+    if cast_type == TypeCastKind.INT_OR_TUPLE_TO_LIST:
+        if isinstance(src_data, list):
+            return src_data
+
+        if isinstance(src_data, int):
+            return [
+                src_data,
+            ]
+
+        if isinstance(src_data, tuple):
+            dst_list = [item for item in src_data]
+            return dst_list
+
+        raise TypeError(f'{src_data} is the wrong data type.')
+
+    raise TypeError("Unsupported type cast")

mindspore/ops/composite/base.py

@@ -20,6 +20,7 @@ from __future__ import absolute_import
 from functools import partial
 
 from types import FunctionType, MethodType
+import numpy as np
 import mindspore as ms
 from mindspore import context
 from mindspore.common.parameter import Parameter, ParameterTuple
@@ -28,7 +29,8 @@ from mindspore._c_expression import GradOperation_, HyperMap_, Map_, MultitypeFu
     TupleAdd_, UnpackCall_, ZipOperation_, ListAppend_, TupleGetItemTensor_, ListInsert_, \
     SequenceSliceGetItem_, ListSliceSetItem_, VmapOperation_, TaylorOperation_, ListPop_, \
     ListClear_, ListReverse_, ListExtend_, DictClear_, DictHasKey_, DictUpdate_, DictFromKeys_, \
-    ZerosLike_, TensorIndexGetitem_, TensorIndexSetitem_, ListAdd_
+    ZerosLike_, TensorIndexGetitem_, TensorIndexSetitem_, ListAdd_, DictSetItem_, \
+    HandleBoolTensor_, HandleEmptySlice_, PreSetitemByTuple_, HandleScalarTensorIndex_
 from mindspore.common import dtype as mstype
 from mindspore.common.api import jit, _pynative_executor, _wrap_func
 from mindspore.common.api import _add_flags, _core
@@ -36,7 +38,8 @@ from mindspore.ops.primitive import Primitive
 from mindspore.ops import signature as sig
 
 __all__ = [TupleAdd_, ListAdd_, UnpackCall_, TupleGetItemTensor_, SequenceSliceGetItem_,
-           ListSliceSetItem_, ZerosLike_, TensorIndexGetitem_, TensorIndexSetitem_]
+           ListSliceSetItem_, ZerosLike_, TensorIndexGetitem_, TensorIndexSetitem_,
+           HandleBoolTensor_, HandleEmptySlice_, PreSetitemByTuple_, HandleScalarTensorIndex_]
 
 
 def add_flags(fn=None, **flags):
@@ -334,7 +337,7 @@ class GradOperation(GradOperation_):
         self.get_all = get_all
         self.get_by_list = get_by_list
         self.sens_param = sens_param
-        GradOperation_.__init__(self, 'grad', get_all, get_by_list, sens_param, False, False, False, False)
+        GradOperation_.__init__(self, 'grad', get_all, get_by_list, sens_param, False, False, False, False, False)
         self.grad_fn = None
         self.fn = None
         self.weights_id = None
@@ -511,8 +514,8 @@ class _Grad(GradOperation_):
     A higher-order function which is used to generate the gradient function by position for the input function.
     """
 
-    def __init__(self, get_by_list=False, sens_param=False, get_by_position=False, has_aux=False, get_value=False,
-                 return_ids=False):
+    def __init__(self, get_all=False, get_by_list=False, sens_param=False, get_by_position=False, has_aux=False,
+                 get_value=False, return_ids=False, merge_forward=False):
         """Initialize _Grad."""
         if not isinstance(get_by_position, bool):
             raise TypeError(f"For '_Grad', the 'get_by_position' should be bool, "
@@ -532,14 +535,16 @@ class _Grad(GradOperation_):
         if not isinstance(return_ids, bool):
             raise TypeError(f"For '_Grad', the 'return_ids' should be bool, "
                             f"but got {type(return_ids).__name__}")
+        self.get_all = get_all
         self.get_by_position = get_by_position
         self.get_by_list = get_by_list
         self.sens_param = sens_param
         self.has_aux = has_aux
         self.get_value = get_value
         self.return_ids = return_ids
-        GradOperation_.__init__(self, 'grad', False, get_by_list, sens_param, get_by_position, has_aux, get_value,
-                                return_ids)
+        self.merge_forward = merge_forward
+        GradOperation_.__init__(self, 'grad', get_all, get_by_list, sens_param, get_by_position, has_aux, get_value,
+                                return_ids, merge_forward)
         self.grad_fn = None
         self.fn = None
         self.pynative_ = False
@@ -562,8 +567,8 @@ class _Grad(GradOperation_):
                 res += (stop_gradient(item),)
             return res
 
-        grad_ = _Grad(self.get_by_list, self.sens_param, self.get_by_position, self.has_aux, self.get_value,
-                      self.return_ids)
+        grad_ = _Grad(self.get_all, self.get_by_list, self.sens_param, self.get_by_position, self.has_aux,
+                      self.get_value, self.return_ids, self.merge_forward)
         # If calling Grad in GRAPH_MODE or calling Grad in functions decorated with 'jit', do grad in GRAPH_MODE
         # If calling Grad in pure PYNATIVE_MODE do grad in PYNATIVE_MODE
         #   In pure PYNATIVE_MODE the out layer after_grad just used to set pynative flag for inner GradOperation.
@@ -738,6 +743,9 @@ class MultitypeFuncGraph(MultitypeFuncGraph_):
                 sig.make_sig('args', sig.sig_rw.RW_READ, sig.sig_kind.KIND_VAR_POSITIONAL),))
 
     def __call__(self, *args):
+        for arg in args:
+            if isinstance(arg, np.ndarray):
+                raise TypeError("For 'MultitypeFuncGraph', the input can not be numpy.ndarray")
         if len(self.entries) == 1:
             output = self.entries[0][1](*args)
             return output
@@ -890,7 +898,7 @@ class Map(Map_):
           If `ops` is `None`, the first input is the operation, and the other is inputs.
 
     Outputs:
-        Sequence, the sequence of output after applying the function. e.g. `operation(args[0][i], args[1][i])`.
+        Sequence, the sequence of output after applying the ops function. e.g. `ops(args[0][i], args[1][i])`.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -1046,6 +1054,25 @@ class _ListExtend(ListExtend_):
 _extend = _ListExtend("extend")
 
 
+class _DictSetItem(DictSetItem_):
+    """
+    A metafuncgraph class that setitem for the dict.
+
+    Args:
+        name (str): The name of the metafuncgraph object.
+    """
+
+    def __init__(self, name):
+        """Initialize _DictClear."""
+        DictSetItem_.__init__(self, name)
+
+    def __call__(self, *args):
+        pass
+
+
+_dict_setitem = _DictSetItem("setitem")
+
+
 class _DictClear(DictClear_):
     """
     A metafuncgraph class that clear the dict.

mindspore/ops/composite/math_ops.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 # ============================================================================
 """math Operations."""
+import mindspore.ops as ops
 from mindspore.ops import functional as F
 from mindspore.ops.function.math_func import cummin as cummin_
+from mindspore.ops._primitive_cache import _get_cache_prim
 
 
 def matmul(x1, x2, dtype=None):
@@ -117,10 +119,9 @@ def mm(input, mat2):
         >>> 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)
+    _matmul = _get_cache_prim(ops.MatMul)()
+    out = _matmul(input, mat2)
+    return out
 
 
 def cummin(x, axis):