mindspore.ops.Primitive

class mindspore.ops.Primitive(name)

Primitive是Python中算子原语的基类。

参数：

• name (str) - 当前Primitive的名称。

样例：

>>> from mindspore.ops import prim_attr_register, Primitive
>>> add = Primitive('add')
>>>
>>> # or work with prim_attr_register:
>>> # init a Primitive class with attr1 and attr2
>>> class Add(Primitive):
...     @prim_attr_register
...     def __init__(self, attr1, attr2):
...         '''init for add'''
...     # check attr1 and attr2 or do some initializations
...     # init a Primitive obj with attr1=1 and attr2=2
>>> add = Add(attr1=1, attr2=2)

add_prim_attr(name, value)

添加Primitive的属性。

参数：

• name (str) - 属性名称。

• value (Any) - 属性值。

样例：

>>> from mindspore import ops
>>> a = ops.Add()
>>> a = a.add_prim_attr("attr",1)
>>> out = a.attrs["attr"]
>>> print(out)
1

check_elim(*args)

检查是否能在编译阶段计算Primitive，如果能，则可以消除。有需要的子类可以重写该方法。

参数：

• args (Primitive args) - 与当前Primitive的参数相同。

返回：

由两个元素组成的元组。第一个元素是指是否能在编译阶段计算Primitive，第二个元素是计算结果。

样例：

>>> import numpy as np
>>> import mindspore
>>> from mindspore import Tensor
>>> from mindspore.ops import prim_attr_register, Primitive
>>> class AddN(Primitive):
...     @prim_attr_register
...     def __init__(self):
...         self.init_prim_io_names(inputs=["inputs"], outputs=["sum"])
...     def check_elim(self, inputs):
...         if len(inputs) != 1:
...             return (False, None)
...         if isinstance(inputs[0], Tensor):
...             return (True, inputs[0])
...
>>> addn = AddN()
>>> input_x = Tensor(np.array([1, 2, 3]), mindspore.float32)
>>> output = addn.check_elim((input_x,))
>>> print(output)
(True, Tensor(shape=[3], dtype=Float32, value= [ 1.00000000e+00,  2.00000000e+00,  3.00000000e+00]))

del_prim_attr(name)

删除Primitive的属性。

参数：

• name (str) - 属性名称。

样例：

>>> from mindspore import ops
>>> a = ops.Add()
>>> a = a.add_prim_attr("attr",1)
>>> a = a.del_prim_attr("attr")
>>> print(a.attrs)
{}

init_prim_io_names(inputs, outputs)

初始化Tensor或属性的输入输出的名称。

参数：

• inputs (list[str]) - 输入名称的列表。

• outputs (list[str]) - 输出名称的列表。

样例：

>>> from mindspore import ops
>>> a = ops.Add()
>>> a.init_prim_io_names(["x","y"],["sum"])
>>> print(a.input_names)
['x','y']
>>> print(a.output_names)
['sum']

place(role, rank_id)

设置Primitive算子标签。让此算子在对应进程上执行。 每个进程的标签都由进程角色 role 和 rank_id 组成，因此，对算子设置标签代表着算子在哪个进程执行，从而执行分布式训练等任务。

说明

• 此接口只在成功调用 mindspore.communication.init() 完成动态组网后才能生效。

参数：

• role (str) - 算子执行所在进程的角色。只支持'MS_WORKER'。

• rank_id (int) - 算子执行所在进程的ID。在相同角色进程间， rank_id 是唯一的。

样例：

>>> from mindspore import context
>>> from mindspore import ops
>>> context.set_context(mode=context.GRAPH_MODE)
>>> matmul = ops.MatMul()
>>> matmul.place('MS_WORKER', 0)

recompute(mode=True)

设置Primitive的重计算属性。

如果有一个被设置了重计算属性的Primitive，并且其结果在计算导数的时候被使用，那么不会保存该Primitive在前向网络中的中间计算结果，而是在自动微分的时候重新进行计算。

说明

• 如果计算涉及随机化或全局变量，则暂无法保证等效性。

• 在PyNative模式下不支持。

参数：

• mode (bool) - Primitive是否设置了重计算。默认值： True 。

样例：

>>> import numpy as np
>>> import mindspore as ms
>>> from mindspore import Tensor, ops, nn, jit
>>> class NetRecompute(nn.Cell):
...     def __init__(self):
...         super(NetRecompute,self).__init__()
...         self.relu = ops.ReLU().recompute()
...         self.sqrt = ops.Sqrt()
...     def construct(self, x):
...         out = self.relu(x)
...         return self.sqrt(out)
...
>>> class GradNet(nn.Cell):
...     def __init__(self, network):
...         super(GradNet,self).__init__()
...         self.network = network
...         self.grad = ops.GradOperation()
...     @jit
...     def construct(self, x):
...         g_out = self.grad(self.network)(x)
...         return g_out
...
>>> x = Tensor(np.array([-1,1]).astype(np.float32))
>>> net = NetRecompute()
>>> grad = GradNet(net)
>>> a = grad(x)
>>> print(a)
[0. 0.5]

set_device(device_target)

设置Primitive执行后端。

参数：

• device_target (str) - 后端名称，支持CPU、GPU、Ascend。

样例：

>>> from mindspore import ops
>>> a = ops.Add()
>>> a = a.set_device("GPU")
>>> print(a.primitive_target)
GPU

set_prim_instance_name(instance_name)

设置Primitive算子的实例的名称。

说明

当用户定义Primitive算子时，默认调用它。

参数：

• instance_name (str) - 用户设置的Primitive算子的实例的名称。

样例：

>>> from mindspore import ops
>>> a = ops.Add()
>>> a = a.set_prim_instance_name("add")
>>> print(a.instance_name)
add

set_stage(stage)

将stage的ID添加到Primitive属性中。

说明

仅在半自动并行模式下有效。在其他并行模式下，请将其设置为0。

参数：

• stage (int) - 当前stage的ID。

样例：

>>> from mindspore import ops
>>> add = ops.Add()
>>> print(add.set_stage(0))
Prim[Add]<stage=0>

shard(in_strategy=None, out_strategy=None)

将切分策略添加到Primitive属性中。

说明

仅在半自动并行或自动并行模式下有效。在其他并行模式中，将忽略此处设置的策略。

参数：

• in_strategy (tuple) - 描述算子输入的切分策略。默认值： None 。

• out_strategy (tuple) - 描述算子输出的切分策略，仅针对某些算子，如MatMul。默认值： None 。

样例：

>>> from mindspore import ops
>>> add = ops.Add()
>>> print(add.shard(((1, 1), (1, 1))))
Prim[Add]<in_strategy=((1, 1), (1, 1)), out_strategy=None>
>>> # using layout
>>> from mindspore import Layout
>>> layout = Layout((2, 2, 2), ("dp", "sp", "mp"))
>>> layout_tuple = (layout("dp", "sp"), layout("sp", "mp"))
>>> from mindspore import ops
>>> matmul = ops.MatMul()
>>> print(matmul.shard(layout_tuple))
Prim[MatMul]<in_layout=({'device_matrix': (2, 2, 2), 'tensor_map': (2, 1)},
{'device_matrix': (2, 2, 2), 'tensor_map': (1, 0)})>
>>> # using layout with None
>>> from mindspore import Layout
>>> layout = Layout((2, 2, 2), ("dp", "sp", "mp"))
>>> layout_tuple = (layout("dp", "sp"), layout("sp", "None")) # "None" means the axis would not be split
>>> from mindspore import ops
>>> matmul = ops.MatMul()
>>> print(matmul.shard(layout_tuple))
Prim[MatMul]<in_layout=({'device_matrix': (2, 2, 2), 'tensor_map': (2, 1)},
{'device_matrix': (2, 2, 2), 'tensor_map': (1, -1)})>

property update_parameter

判断此Primitive是否会更新参数的值。