mindspore 1.10.0__cp37-cp37m-win_amd64.whl → 2.0.0rc1__cp37-cp37m-win_amd64.whl

mindspore/nn/layer/rnn_cells.py

@@ -24,9 +24,9 @@ from mindspore import log as logger
 from mindspore.common.tensor import Tensor
 from mindspore.common.parameter import Parameter
 from mindspore.common.initializer import initializer, Uniform
-from mindspore.ops.primitive import constexpr
+from mindspore.ops.primitive import constexpr, _primexpr
 from mindspore.nn.cell import Cell
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 
 __all__ = ['LSTMCell', 'GRUCell', 'RNNCell']
 
@@ -36,7 +36,7 @@ def _check_input_dtype(input_dtype, param_name, allow_dtypes, cls_name):
     validator.check_type_name(param_name, input_dtype, allow_dtypes, cls_name)
 
 
-@constexpr
+@constexpr(check=False)
 def _check_is_tensor(param_name, input_data, cls_name):
     """Internal function, used to check whether the input data is Tensor."""
     if input_data is not None and not isinstance(P.typeof(input_data), mstype.tensor_type):
@@ -60,7 +60,7 @@ def _check_tuple_length(param_name, input_data, length, cls_name):
                         f"but got '{len(input_data)}'")
 
 
-@constexpr
+@_primexpr
 def _check_batch_size_equal(batch_size_x, batch_size_hx, cls_name):
     if batch_size_x != batch_size_hx:
         raise ValueError(f"For '{cls_name}' batch size of x and hx must be equal, but got {batch_size_x} of x "
@@ -175,7 +175,7 @@ class RNNCell(RNNCellBase):
 
     Here :math:`h_t` is the hidden state at time `t`, :math:`x_t` is
     the input at time `t`, and :math:`h_{(t-1)}` is the hidden state of the
-    previous layer at time `t-1` or the initial hidden state at time `0`.
+    previous layer at time :math:`t-1` or the initial hidden state at time `0`.
     If `nonlinearity` is `relu`, then `relu` is used instead of `tanh`.
 
     Args:
@@ -266,12 +266,12 @@ class LSTMCell(RNNCellBase):
         has_bias (bool): Whether the cell has bias `b_ih` and `b_hh`. Default: True.
 
     Inputs:
-        - **x** (Tensor) - Tensor of shape (batch_size, `input_size`).
+        - **x** (Tensor) - Tensor of shape :math:`(batch\_size, input\_size)`.
         - **hx** (tuple) - A tuple of two Tensors (h_0, c_0) both of data type mindspore.float32
-          and shape (batch_size, `hidden_size`). The data type of `hx` must be the same as `x`.
+          and shape :math:`(batch\_size, hidden\_size)`. The data type of `hx` must be the same as `x`.
 
     Outputs:
-        - **hx'** (Tensor) - A tuple of two Tensors (h', c') both of data shape (batch_size, `hidden_size`).
+        - **hx'** (Tensor) - A tuple of two Tensors (h', c') both of data shape :math:`(batch\_size, hidden\_size)`.
 
     Raises:
         TypeError: If `input_size`, `hidden_size` is not an int.
@@ -340,23 +340,18 @@ class GRUCell(RNNCellBase):
     `Learning Phrase Representations using RNN Encoder–Decoder for Statistical Machine Translation
     <https://aclanthology.org/D14-1179.pdf>`_.
 
-    The LSTMCell can be simplified in NN layer, the following formula:
-
-    .. math::
-        h^{'},c^{'} = LSTMCell(x, (h_0, c_0))
-
     Args:
         input_size (int): Number of features of input.
         hidden_size (int):  Number of features of hidden layer.
         has_bias (bool): Whether the cell has bias `b_in` and `b_hn`. Default: True.
 
     Inputs:
-        - **x** (Tensor) - Tensor of shape (batch_size, `input_size`).
-        - **hx** (Tensor) - Tensor of data type mindspore.float32 and shape (batch_size, `hidden_size`).
+        - **x** (Tensor) - Tensor of shape :math:`(batch\_size, input\_size)`.
+        - **hx** (Tensor) - Tensor of data type mindspore.float32 and shape :math:`(batch\_size, hidden\_size)`.
           Data type of `hx` must be the same as `x`.
 
     Outputs:
-        - **hx'** (Tensor) - Tensor of shape (batch_size, `hidden_size`).
+        - **hx'** (Tensor) - Tensor of shape :math:`(batch\_size, hidden\_size)`.
 
     Raises:
         TypeError: If `input_size`, `hidden_size` is not an int.

mindspore/nn/layer/rnns.py

@@ -22,28 +22,23 @@ import mindspore.nn as nn
 import mindspore.ops as P
 import mindspore.context as context
 import mindspore.common.dtype as mstype
-from mindspore.ops.primitive import constexpr
+from mindspore.ops import functional as F
+from mindspore.ops.primitive import constexpr, _primexpr
 from mindspore.common.tensor import Tensor
 from mindspore.common.parameter import ParameterTuple, Parameter
 from mindspore.nn.cell import Cell
 from mindspore import log as logger
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.ops.operations._rl_inner_ops import CudnnGRU
 from mindspore.nn.layer.rnn_cells import _rnn_relu_cell, _rnn_tanh_cell, _gru_cell, _lstm_cell
-from mindspore.nn.layer.rnn_utils import _Reverse, _ReverseSequence
 
 __all__ = ['LSTM', 'GRU', 'RNN']
 
 
-@constexpr
-def arange(start, stop, step, dtype):
-    return Tensor(np.arange(start, stop, step), dtype)
-
-
-@constexpr
+@_primexpr
 def _init_state(shape, dtype, is_lstm):
-    hx = Tensor(np.zeros(shape), dtype)
-    cx = Tensor(np.zeros(shape), dtype)
+    hx = P.Zeros()(shape, dtype)
+    cx = P.Zeros()(shape, dtype)
     if is_lstm:
         return (hx, cx)
     return hx
@@ -60,7 +55,7 @@ def _check_input_dtype_same_and_valid(args_name, args_value, valid_values, cls_n
     validator.check_types_same_and_valid(args, valid_values, cls_name)
 
 
-@constexpr
+@constexpr(check=False)
 def _check_is_tensor(param_name, input_data, cls_name):
     """Internal function, used to check whether the input data is Tensor."""
     if input_data is not None and not isinstance(P.typeof(input_data), mstype.tensor_type):
@@ -84,7 +79,7 @@ def _check_tuple_length(param_name, input_data, length, cls_name):
                         f"but got '{len(input_data)}'")
 
 
-@constexpr
+@_primexpr
 def _check_seq_length_size(batch_size_x, seq_length_size, cls_name):
     if batch_size_x != seq_length_size:
         raise ValueError(f"For '{cls_name}' batch size of x and seq_length must be equal, "
@@ -93,7 +88,7 @@ def _check_seq_length_size(batch_size_x, seq_length_size, cls_name):
 
 def sequence_mask(lengths, maxlen):
     """generate mask matrix by seq_length"""
-    range_vector = arange(0, maxlen, 1, lengths.dtype)
+    range_vector = P.arange(start=0, end=maxlen, step=1, dtype=lengths.dtype)
     result = range_vector < lengths.view(lengths.shape + (1,))
     return result.astype(mstype.int32)
 
@@ -106,7 +101,7 @@ def select_by_mask(inputs, mask):
 
 def get_hidden(output, seq_length):
     """get hidden state by seq_length"""
-    batch_index = arange(0, seq_length.shape[0], 1, seq_length.dtype)
+    batch_index = P.arange(start=0, end=seq_length.shape[0], step=1, dtype=seq_length.dtype)
     indices = P.Concat(1)((seq_length.view(-1, 1) - 1, batch_index.view(-1, 1)))
     return P.GatherNd()(output, indices)
 
@@ -158,7 +153,7 @@ class _DynamicRNNBase(Cell):
             hidden_size = h.shape[-1]
             zero_output = P.ZerosLike()(h_t)
         seq_length = P.Cast()(seq_length, mstype.float32)
-        seq_length = P.BroadcastTo((hidden_size, -1))(seq_length)
+        seq_length = F.broadcast_to(seq_length, (hidden_size, -1))
         seq_length = P.Cast()(seq_length, mstype.int32)
         seq_length = P.Transpose()(seq_length, (1, 0))
 
@@ -220,6 +215,7 @@ class _DynamicGRUCPUGPU(Cell):
         self.is_gpu = context.get_context("device_target") == "GPU"
 
     def construct(self, x, h_0, seq_length, w_ih, w_hh, b_ih, b_hh):
+        '''_DynamicGRUCPUGPU'''
         gate_size, input_size = w_ih.shape
         hidden_size = gate_size // 3
         if self.is_gpu:
@@ -262,15 +258,16 @@ class _DynamicGRUAscend(Cell):
         self.dtype = mstype.float16
 
     def construct(self, x, h_0, seq_length, w_ih, w_hh, b_ih, b_hh):
+        '''Dynamic GRU module on Ascend'''
         if b_ih is None:
             b_ih = P.Zeros()(w_ih.shape[0], w_ih.dtype)
             b_hh = P.Zeros()(w_ih.shape[0], w_ih.dtype)
         outputs, _, _, _, _, _ = self.gru(self.cast(x, self.dtype), \
                                           self.cast(self.transpose(w_ih, (1, 0)), self.dtype), \
                                           self.cast(self.transpose(w_hh, (1, 0)), self.dtype), \
-                                          self.cast(b_ih, self.dtype), \
-                                          self.cast(b_hh, self.dtype), \
-                                          None, self.cast(h_0, self.dtype))
+                                          b_ih, \
+                                          b_hh, \
+                                          None, h_0)
         if seq_length is not None:
             h = get_hidden(outputs, seq_length)
             mask = sequence_mask(seq_length, x.shape[0])
@@ -289,6 +286,7 @@ class _DynamicLSTMCPUGPU(Cell):
         self.is_gpu = context.get_context("device_target") == "GPU"
 
     def construct(self, x, h_0, seq_length, w_ih, w_hh, b_ih, b_hh):
+        '''Dynamic LSTM module on CPU and GPU'''
         gate_size, input_size = w_ih.shape
         hidden_size = gate_size // 4
         if seq_length is not None:
@@ -339,6 +337,7 @@ class _DynamicLSTMAscend(Cell):
         self.dtype = mstype.float16
 
     def construct(self, x, h_0, seq_length, w_ih, w_hh, b_ih, b_hh):
+        '''Dynamic LSTM module on Ascend'''
         w_ih_i, w_ih_f, w_ih_g, w_ih_o = self.split(w_ih)
         w_hh_i, w_hh_f, w_hh_g, w_hh_o = self.split(w_hh)
         w_ih = self.concat_dim0((w_ih_i, w_ih_g, w_ih_f, w_ih_o))
@@ -414,17 +413,13 @@ class _RNNBase(Cell):
             raise ValueError(f"For '{self.cls_name}', the 'mode' must be in ['RNN_RELU', 'RNN_TANH', 'LSTM', 'GRU'], "
                              f"but got {mode}.")
 
-        if context.get_context("device_target") == "CPU":
-            self.reverse = _Reverse(0)
-            self.reverse_sequence = _ReverseSequence(0, 1)
-        else:
-            self.reverse = P.ReverseV2([0])
-            self.reverse_sequence = P.ReverseSequence(0, 1)
+        self.reverse = P.ReverseV2([0])
+        self.reverse_sequence = P.ReverseSequence(0, 1)
         self.hidden_size = hidden_size
         self.batch_first = batch_first
         self.num_layers = num_layers
         self.dropout = dropout
-        self.dropout_op = nn.Dropout(float(1 - dropout))
+        self.dropout_op = nn.Dropout(p=float(dropout))
         self.bidirectional = bidirectional
         self.has_bias = has_bias
         num_directions = 2 if bidirectional else 1
@@ -503,9 +498,11 @@ class _RNNBase(Cell):
         if self.is_lstm:
             h_n = P.Concat(0)(h_n)
             c_n = P.Concat(0)(c_n)
-            h_n = h_n.view(h[0].shape)
-            c_n = c_n.view(h[1].shape)
-            return output, (h_n.view(h[0].shape), c_n.view(h[1].shape))
+            h0_shape = h[0].shape
+            h1_shape = h[1].shape
+            h_n = h_n.view(h0_shape)
+            c_n = c_n.view(h1_shape)
+            return output, (h_n.view(h0_shape), c_n.view(h1_shape))
         h_n = P.Concat(0)(h_n)
         return output, h_n.view(h.shape)
 
@@ -535,9 +532,11 @@ class _RNNBase(Cell):
         if self.is_lstm:
             h_n = P.Concat(0)(h_n)
             c_n = P.Concat(0)(c_n)
-            h_n = h_n.view(h[0].shape)
-            c_n = c_n.view(h[1].shape)
-            return output, (h_n.view(h[0].shape), c_n.view(h[1].shape))
+            h0_shape = h[0].shape
+            h1_shape = h[1].shape
+            h_n = h_n.view(h0_shape)
+            c_n = c_n.view(h1_shape)
+            return output, (h_n.view(h0_shape), c_n.view(h1_shape))
         h_n = P.Concat(0)(h_n)
         return output, h_n.view(h.shape)
 
@@ -591,7 +590,7 @@ class RNN(_RNNBase):
 
     Here :math:`h_t` is the hidden state at time `t`, :math:`x_t` is
     the input at time `t`, and :math:`h_{(t-1)}` is the hidden state of the
-    previous layer at time `t-1` or the initial hidden state at time `0`.
+    previous layer at time :math:`t-1` or the initial hidden state at time `0`.
     If :attr:`nonlinearity` is ``'relu'``, then :math:`\text{ReLU}` is used instead of :math:`\tanh`.
 
     Args:
@@ -671,7 +670,7 @@ class GRU(_RNNBase):
     Given an input :math:`x_t` at time :math:`t`, a hidden state :math:`h_{t-1}`, the update and reset gate at
     time :math:`t` is computed using a gating mechanism. Update gate :math:`z_t` is designed to protect the cell
     from perturbation by irrelevant inputs and past hidden state. Reset gate :math:`r_t` determines how much
-    information should be reset from old hidden state. New memory state :math:`{n}_t` is
+    information should be reset from old hidden state. New memory state :math:`n_t` is
     calculated with the current input, on which the reset gate will be applied. Finally, current hidden state
     :math:`h_{t}` is computed with the calculated update grate and new memory state. The complete
     formulation is as follows:
@@ -805,12 +804,12 @@ class LSTM(_RNNBase):
 
     Inputs:
         - **x** (Tensor) - Tensor of data type mindspore.float32 or mindspore.float16 and
-          shape (seq_len, batch_size, `input_size`) or (batch_size, seq_len, `input_size`).
+          shape :math:`(seq\_len, batch\_size, input\_size)` or :math:`(batch\_size, seq\_len, input\_size)`.
         - **hx** (tuple) - A tuple of two Tensors (h_0, c_0) both of data type mindspore.float32
-          or mindspore.float16 and shape (num_directions * `num_layers`, batch_size, `hidden_size`).
+          or mindspore.float16 and shape :math:`(num\_directions * num\_layers, batch\_size, hidden\_size)`.
           The data type of `hx` must be the same as `x`.
         - **seq_length** (Tensor) - The length of each sequence in an input batch.
-          Tensor of shape :math:`(\text{batch_size})`. Default: None.
+          Tensor of shape :math:`(batch\_size)`. Default: None.
           This input indicates the real sequence length before padding to avoid padded elements
           have been used to compute hidden state and affect the final output. It is recommended to
           use this input when **x** has padding elements.
@@ -818,9 +817,9 @@ class LSTM(_RNNBase):
     Outputs:
         Tuple, a tuple contains (`output`, (`h_n`, `c_n`)).
 
-        - **output** (Tensor) - Tensor of shape (seq_len, batch_size, num_directions * `hidden_size`).
+        - **output** (Tensor) - Tensor of shape :math:`(seq\_len, batch\_size, num\_directions * hidden\_size)` .
         - **hx_n** (tuple) - A tuple of two Tensor (h_n, c_n) both of shape
-          (num_directions * `num_layers`, batch_size, `hidden_size`).
+          :math:`(num\_directions * num\_layers, batch\_size, hidden\_size)` .
 
     Raises:
         TypeError: If `input_size`, `hidden_size` or `num_layers` is not an int.

mindspore/nn/layer/thor_layer.py

@@ -25,7 +25,8 @@ from mindspore.communication.management import get_group_size, get_rank
 from mindspore.ops import operations as P
 from mindspore.ops.operations._thor_ops import ThorIm2Col
 from mindspore.common.parameter import Parameter
-from mindspore._checkparam import Validator, Rel, twice
+from mindspore import _checkparam as Validator
+from mindspore._checkparam import twice
 from mindspore import context
 from mindspore.nn.cell import Cell
 from mindspore.nn.layer.activation import get_activation
@@ -33,9 +34,9 @@ from mindspore.parallel._ps_context import _is_role_worker, _get_ps_context, \
     _set_rank_id, _insert_hash_table_size, _set_cache_enable
 from mindspore.parallel._utils import _get_parallel_mode, _get_full_batch
 from mindspore.context import ParallelMode
-from mindspore.ops.primitive import constexpr
 from mindspore.ops import functional as F
 from mindspore.nn.layer.basic import ClipByNorm
+from mindspore.ops.primitive import constexpr
 
 __all__ = ['DenseThor', 'Conv2dThor', 'EmbeddingThor', 'EmbeddingLookupThor']
 
@@ -45,8 +46,8 @@ class DenseThor(Cell):
     The dense connected layer and saving the information needed for THOR.
 
     Applies dense connected layer for the input and saves the information A and G in the dense connected layer
-    needed for THOR, the detail can be seen in `THOR, Trace-based Hardware-driven layer-ORiented Natural
-    Gradient Descent Computation <https://www.aaai.org/AAAI21Papers/AAAI-6611.ChenM.pdf>`_.
+    needed for THOR.
+
     This layer implements the operation as:
 
     .. math::
@@ -131,12 +132,12 @@ class DenseThor(Cell):
         self.reshape = P.Reshape()
         self.transpose = P.Transpose()
         self.mul = P.Mul()
-        self.is_Ascend = True
+        self.is_ascend = True
         self.split_dim = 128
         if context.get_context("device_target") == "Ascend":
             self._process_ascend_dense_thor(out_channels, in_channels)
         else:
-            self.is_Ascend = False
+            self.is_ascend = False
             self.cube_matmul = P.MatMul(transpose_a=True)
         self.getG = P.InsertGradientOf(self.save_gradient)
 
@@ -153,7 +154,7 @@ class DenseThor(Cell):
            save_gradient
         """
         out = dout
-        if self.is_Ascend:
+        if self.is_ascend:
             if not self.is_nsp_layer:
                 shape = self.shape(dout)
                 normalizer = self.cast(shape[0], mstype.float32)
@@ -170,7 +171,7 @@ class DenseThor(Cell):
 
     def construct(self, x):
         if self.thor:
-            if self.is_Ascend:
+            if self.is_ascend:
                 inputs = self.cube_matmul(x, x)
                 shape = self.shape(x)
                 normalizer = self.cast(shape[0], mstype.float32)
@@ -283,7 +284,6 @@ class Conv2dThor(_ConvThor):
     Applies a 2D convolution over an input tensor which is typically of shape :math:`(N, C_{in}, H_{in}, W_{in})`,
     where :math:`N` is batch size, :math:`C_{in}` is channel number, and :math:`H_{in}, W_{in})` are height and width.
     And saves the information A and G in the 2D convolution layer needed for THOR.
-    The detail can be seen in paper: https://www.aaai.org/AAAI21Papers/AAAI-6611.ChenM.pdf
 
     For each batch of shape :math:`(C_{in}, H_{in}, W_{in})`, the formula is defined as:
 
@@ -399,11 +399,11 @@ class Conv2dThor(_ConvThor):
         self.cast = P.Cast()
         self.a_normalizer = Parameter(initializer(1, [1], mstype.float32), name="a_normalizer", requires_grad=False)
         self.g_normalizer = Parameter(initializer(1, [1], mstype.float32), name="g_normalizer", requires_grad=False)
-        self.is_Ascend = True
+        self.is_ascend = True
         if context.get_context("device_target") == "Ascend":
             self._process_ascend_conv2d_thor(kernel_size, stride)
         else:
-            self.is_Ascend = False
+            self.is_ascend = False
             self.img2col = ThorIm2Col(kernel_size=kernel_size, stride=stride, pad_mode="same")
             self.matmul = P.MatMul(transpose_b=True)
             self.reduce_mean = P.ReduceMean(keep_dims=False)
@@ -434,8 +434,8 @@ class Conv2dThor(_ConvThor):
         """Initialize depthwise conv2d op"""
         if context.get_context("device_target") == "Ascend" and self.group > 1:
             self.dilation = self._dilation
-            Validator.check_int('group', self.group, self.in_channels, Rel.EQ, self.cls_name)
-            Validator.check_int('group', self.group, self.out_channels, Rel.EQ, self.cls_name)
+            Validator.check_int('group', self.group, self.in_channels, Validator.EQ, self.cls_name)
+            Validator.check_int('group', self.group, self.out_channels, Validator.EQ, self.cls_name)
             self.conv2d = P.DepthwiseConv2dNative(channel_multiplier=1,
                                                   kernel_size=self.kernel_size,
                                                   pad_mode=self.pad_mode,
@@ -453,7 +453,7 @@ class Conv2dThor(_ConvThor):
     def save_gradient(self, dout):
         """save_gradient"""
         out = dout
-        if self.is_Ascend:
+        if self.is_ascend:
             dout_shape = self.shape(dout)
             dout = self.transpose(dout, (0, 2, 3, 1))
             dout = self.reshape(dout, (-1, dout_shape[1]))
@@ -479,7 +479,7 @@ class Conv2dThor(_ConvThor):
 
     def construct(self, x):
         if self.thor:
-            if self.is_Ascend:
+            if self.is_ascend:
                 matrix_a = self.im2col(x)
                 matrix_a_shape = self.shape(matrix_a)
                 y = matrix_a_shape[3]
@@ -510,13 +510,13 @@ class Conv2dThor(_ConvThor):
                 output = self.conv2d(x, self.weight)
                 output = self.getG(output)
         else:
-            if self.is_Ascend:
+            if self.is_ascend:
                 weight = self.cast(self.weight, mstype.float16)
                 output = self.conv2d(x, weight)
             else:
                 output = self.conv2d(x, self.weight)
         if self.has_bias:
-            if self.is_Ascend:
+            if self.is_ascend:
                 bias = self.cast(self.bias, mstype.float16)
                 output = self.bias_add(output, bias)
             else:
@@ -540,7 +540,7 @@ class EmbeddingThor(Cell):
     This module is often used to store word embeddings and retrieve them using
     indices. The input to the module is a list of indices, and the output is
     the corresponding word embeddings. And saves the information A and G in the dense connected layer
-    needed for THOR, the detail can be seen in paper: https://www.aaai.org/AAAI21Papers/AAAI-6611.ChenM.pdf
+    needed for THOR.
 
     Note:
         When 'use_one_hot' is set to True, the type of the input `x` must be mindspore.int32.
@@ -588,9 +588,9 @@ class EmbeddingThor(Cell):
         self.init_tensor = initializer(embedding_table, [vocab_size, embedding_size])
         self.padding_idx = padding_idx
         if padding_idx is not None:
-            self.padding_idx = Validator.check_int_range(padding_idx, 0, vocab_size, Rel.INC_BOTH,
+            self.padding_idx = Validator.check_int_range(padding_idx, 0, vocab_size, Validator.INC_BOTH,
                                                          "padding_idx", self.cls_name)
-            self.init_tensor = self.init_tensor.to_tensor().asnumpy()
+            self.init_tensor = self.init_tensor.init_data().asnumpy()
             self.init_tensor[self.padding_idx] = 0
         self.embedding_table = Parameter(self.init_tensor, name='embedding_table')
         self.expand = P.ExpandDims()
@@ -671,8 +671,7 @@ class EmbeddingLookupThor(Cell):
     and saving the information needed for THOR.
 
     This module has the same function as EmbeddingLookup, but additionally saves the information A and G in the
-    embeddinglookup layer needed for THOR,
-    the detail can be seen in paper: https://www.aaai.org/AAAI21Papers/AAAI-6611.ChenM.pdf
+    embeddinglookup layer needed for THOR.
 
 
     Args:
@@ -749,7 +748,7 @@ class EmbeddingLookupThor(Cell):
             self.gatherv2 = P.SparseGatherV2()
         else:
             self.gatherv2 = P.Gather()
-        self.embeddinglookup = P.EmbeddingLookup().add_prim_attr('primitive_target', 'CPU')
+        self.embeddinglookup = P.EmbeddingLookup().set_device('CPU')
         enable_ps = _get_ps_context("enable_ps")
         if enable_ps:
             self._process_vocab_cache(slice_mode)
@@ -866,11 +865,11 @@ class EmbeddingLookupThor(Cell):
 
         logger.info("EmbeddingLookup cache enable takes effect.")
         self.forward_unique = True
-        self.unique = P.Unique().add_prim_attr('primitive_target', 'CPU')
+        self.unique = P.Unique().set_device('CPU')
         self.unique.add_prim_attr('cache_enable', True)
         self.embedding_table.cache_enable = self.cache_enable
         self.embedding_table.cache_shape = (self.vocab_cache_size, self.embedding_size)
-        self.reshape_first = P.Reshape().add_prim_attr('primitive_target', 'CPU')
+        self.reshape_first = P.Reshape().set_device('CPU')
 
     def _process_vocab_cache(self, slice_mode):
         """PS embeddingLookup cache check and process."""

mindspore/nn/layer/timedistributed.py

@@ -15,16 +15,16 @@
 """Time Distributed."""
 from __future__ import absolute_import
 
-from mindspore.ops.primitive import constexpr, Primitive
+from mindspore.ops.primitive import constexpr, Primitive, _primexpr
 from mindspore.ops import Reshape, Transpose, Stack, Unstack
 from mindspore.common import Tensor
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.nn.cell import Cell
 
 __all__ = ['TimeDistributed']
 
 
-@constexpr
+@_primexpr
 def _check_reshape_pos(reshape_pos, inputs_shape, outputs_shape, prim_name=None):
     msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
     if reshape_pos >= len(outputs_shape) or inputs_shape[reshape_pos] != outputs_shape[reshape_pos]:
@@ -35,7 +35,7 @@ def _check_reshape_pos(reshape_pos, inputs_shape, outputs_shape, prim_name=None)
                          f"{outputs_shape}. You may try pass parameters without 'reshape_with_axis'.")
 
 
-@constexpr
+@_primexpr
 def _check_expand_dims_axis(time_axis, ndim, prim_name=None):
     msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
     if time_axis > ndim:
@@ -57,7 +57,7 @@ def _check_data(flag, prim_name=None):
         raise TypeError(f"{msg_prefix} inputs and outputs must be a Tensor.")
 
 
-@constexpr
+@_primexpr
 def _check_inputs_dim(shape, prim_name=None):
     msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
     if len(shape) < 3: