diff --git a/mindspore/ops/operations/__init__.py b/mindspore/ops/operations/__init__.py index 4c2f76aab1..8fcf364402 100644 --- a/mindspore/ops/operations/__init__.py +++ b/mindspore/ops/operations/__init__.py @@ -41,7 +41,7 @@ from .comm_ops import (AllGather, AllReduce, _AlltoAll, AllSwap, ReduceScatter, from .debug_ops import (ImageSummary, InsertGradientOf, HookBackward, ScalarSummary, TensorSummary, HistogramSummary, Print, Assert) from .control_ops import ControlDepend, GeSwitch, Merge -from .inner_ops import ScalarCast, Randperm, NoRepeatNGram, LambApplyOptimizerAssign, LambApplyWeightAssign +from .inner_ops import ScalarCast, Randperm, NoRepeatNGram, LambApplyOptimizerAssign, LambApplyWeightAssign, MakeRefKey from .math_ops import (Abs, ACos, Asin, Asinh, AddN, AccumulateNV2, AssignAdd, AssignSub, Atan2, BatchMatMul, BitwiseAnd, BitwiseOr, BitwiseXor, Inv, Invert, ApproximateEqual, InplaceAdd, InplaceSub, @@ -86,7 +86,7 @@ from . import _quant_ops from ._quant_ops import * from .other_ops import (Assign, InplaceAssign, IOU, BoundingBoxDecode, BoundingBoxEncode, ConfusionMatrix, PopulationCount, - CheckValid, MakeRefKey, Partial, Depend, identity, CheckBprop, Push, Pull) + CheckValid, Partial, Depend, identity, CheckBprop, Push, Pull) from ._thor_ops import (CusBatchMatMul, CusCholeskyTrsm, CusFusedAbsMax1, CusImg2Col, CusMatMulCubeDenseLeft, CusMatMulCubeFraczRightMul, CusMatMulCube, CusMatrixCombine, CusTranspose02314, CusMatMulCubeDenseRight, @@ -290,8 +290,8 @@ __all__ = [ 'Floor', 'NMSWithMask', 'IOU', - 'MakeRefKey', 'Partial', + 'MakeRefKey', 'Depend', 'identity', 'AvgPool', diff --git a/mindspore/ops/operations/_inner_ops.py b/mindspore/ops/operations/_inner_ops.py index 60f28d108b..1b3ce2a731 100644 --- a/mindspore/ops/operations/_inner_ops.py +++ b/mindspore/ops/operations/_inner_ops.py @@ -435,7 +435,7 @@ class Receive(PrimitiveWithInfer): will be send by the Send op with the same "sr_tag". src_rank (int): A required integer identifying the source rank. shape (list[int]): A required list identifying the shape of the tensor to be received. - dtype (Type): A required Type indentifying the type of the tensor to be received. The supported types: + dtype (Type): A required Type identifying the type of the tensor to be received. The supported types: int8, int16, int32, float16, float32. group (str): The communication group to work on. Default: "hccl_world_group/nccl_world_group". diff --git a/mindspore/ops/operations/inner_ops.py b/mindspore/ops/operations/inner_ops.py index 99eabe83cc..efd5ef3ff1 100644 --- a/mindspore/ops/operations/inner_ops.py +++ b/mindspore/ops/operations/inner_ops.py @@ -20,7 +20,7 @@ from ..._checkparam import Validator as validator from ..._checkparam import Rel from ...common import dtype as mstype from ...common.dtype import tensor, dtype_to_pytype -from ..primitive import prim_attr_register, PrimitiveWithInfer +from ..primitive import prim_attr_register, Primitive, PrimitiveWithInfer class ScalarCast(PrimitiveWithInfer): @@ -308,3 +308,52 @@ class LambApplyWeightAssign(PrimitiveWithInfer): args = {"w_norm": w_norm_dtype, "g_norm": g_norm_dtype, "lr": lr_dtype} validator.check_scalar_or_tensor_types_same(args, [mstype.float16, mstype.float32], self.name, True) return var_dtype + + +class MakeRefKey(Primitive): + """ + Makes a RefKey instance by string. RefKey stores the name of Parameter, can be passed through the functions, + and used for Assign target. + + Args: + tag (str): Parameter name to make the RefKey. + + Inputs: + No inputs. + + Outputs: + RefKeyType, made from the Parameter name. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + + Examples: + >>> import numpy as np + >>> from mindspore import Parameter, Tensor + >>> from mindspore import dtype as mstype + >>> import mindspore.ops as ops + >>> class Net(nn.Cell): + ... def __init__(self): + ... super(Net, self).__init__() + ... self.y = Parameter(Tensor(np.ones([2, 3]), mstype.int32), name="y") + ... self.make_ref_key = ops.MakeRefKey("y") + ... + ... def construct(self, x): + ... key = self.make_ref_key() + ... ref = ops.make_ref(key, x, self.y) + ... return ref * x + ... + >>> x = Tensor(np.array([[1, 2, 3], [4, 5, 6]]), mindspore.int32) + >>> net = Net() + >>> output = net(x) + >>> print(output) + [[ 1 4 9] + [16 25 36]] + """ + + @prim_attr_register + def __init__(self, tag): + validator.check_value_type('tag', tag, (str,), self.name) + + def __call__(self): + pass diff --git a/mindspore/ops/operations/other_ops.py b/mindspore/ops/operations/other_ops.py index f87d684000..e21b3b535c 100644 --- a/mindspore/ops/operations/other_ops.py +++ b/mindspore/ops/operations/other_ops.py @@ -343,55 +343,6 @@ class IOU(PrimitiveWithInfer): return anchor_boxes -class MakeRefKey(Primitive): - """ - Makes a RefKey instance by string. RefKey stores the name of Parameter, can be passed through the functions, - and used for Assign target. - - Args: - tag (str): Parameter name to make the RefKey. - - Inputs: - No inputs. - - Outputs: - RefKeyType, made from the Parameter name. - - Supported Platforms: - ``Ascend`` ``GPU`` ``CPU`` - - Examples: - >>> import numpy as np - >>> from mindspore import Parameter, Tensor - >>> from mindspore import dtype as mstype - >>> import mindspore.ops as ops - >>> class Net(nn.Cell): - ... def __init__(self): - ... super(Net, self).__init__() - ... self.y = Parameter(Tensor(np.ones([2, 3]), mstype.int32), name="y") - ... self.make_ref_key = ops.MakeRefKey("y") - ... - ... def construct(self, x): - ... key = self.make_ref_key() - ... ref = ops.make_ref(key, x, self.y) - ... return ref * x - ... - >>> x = Tensor(np.array([[1, 2, 3], [4, 5, 6]]), mindspore.int32) - >>> net = Net() - >>> output = net(x) - >>> print(output) - [[ 1 4 9] - [16 25 36]] - """ - - @prim_attr_register - def __init__(self, tag): - validator.check_value_type('tag', tag, (str,), self.name) - - def __call__(self): - pass - - class Partial(Primitive): """ Makes a partial function instance, used for pynative mode.