mindspore 2.0.0rc1__cp38-cp38-manylinux1_x86_64.whl → 2.2.0__cp38-cp38-manylinux1_x86_64.whl

mindspore/nn/layer/rnn_cells.py

@@ -14,6 +14,7 @@
 # ============================================================================
 """RNN Cells module, include RNNCell, GRUCell, LSTMCell."""
 from __future__ import absolute_import
+from functools import wraps
 
 import math
 import numpy as np
@@ -39,8 +40,8 @@ def _check_input_dtype(input_dtype, param_name, allow_dtypes, cls_name):
 @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):
-        raise TypeError(f"For '{cls_name}', the '{param_name}' must be '{mstype.tensor_type}', "
+    if input_data is not None and not isinstance(P.typeof(input_data), mstype.TensorType):
+        raise TypeError(f"For '{cls_name}', the '{param_name}' must be '{mstype.TensorType}', "
                         f"but got '{P.typeof(input_data)}'")
 
 
@@ -68,6 +69,8 @@ def _check_batch_size_equal(batch_size_x, batch_size_hx, cls_name):
 
 
 def _check_lstmcell_init(func):
+    """Internal function, used to check init args."""
+    @wraps(func)
     def wrapper(*args, **kwargs):
         logger.warning(f"LSTMCell has been changed from 'single LSTM layer' to 'single LSTM cell', "
                        f"if you still need use single LSTM layer, please use `nn.LSTM` instead.")
@@ -142,7 +145,8 @@ def _gru_cell(inputs, hidden, w_ih, w_hh, b_ih, b_hh):
 
 class RNNCellBase(Cell):
     '''Basic class for RNN Cells'''
-    def __init__(self, input_size: int, hidden_size: int, has_bias: bool, num_chunks: int):
+    def __init__(self, input_size: int, hidden_size: int, has_bias: bool, num_chunks: int,
+                 dtype=mstype.float32):
         super().__init__()
         validator.check_value_type("has_bias", has_bias, [bool], self.cls_name)
         validator.check_positive_int(hidden_size, "hidden_size", self.cls_name)
@@ -150,20 +154,20 @@ class RNNCellBase(Cell):
         self.input_size = input_size
         self.hidden_size = hidden_size
         self.has_bias = has_bias
-        self.weight_ih = Parameter(Tensor(np.random.randn(num_chunks * hidden_size, input_size).astype(np.float32)))
-        self.weight_hh = Parameter(Tensor(np.random.randn(num_chunks * hidden_size, hidden_size).astype(np.float32)))
+        self.weight_ih = Parameter(Tensor(np.random.randn(num_chunks * hidden_size, input_size), dtype=dtype))
+        self.weight_hh = Parameter(Tensor(np.random.randn(num_chunks * hidden_size, hidden_size), dtype=dtype))
         if has_bias:
-            self.bias_ih = Parameter(Tensor(np.random.randn(num_chunks * hidden_size).astype(np.float32)))
-            self.bias_hh = Parameter(Tensor(np.random.randn(num_chunks * hidden_size).astype(np.float32)))
+            self.bias_ih = Parameter(Tensor(np.random.randn(num_chunks * hidden_size), dtype=dtype))
+            self.bias_hh = Parameter(Tensor(np.random.randn(num_chunks * hidden_size), dtype=dtype))
         else:
             self.bias_ih = None
             self.bias_hh = None
-        self.reset_parameters()
+        self.reset_parameters(dtype=dtype)
 
-    def reset_parameters(self):
+    def reset_parameters(self, dtype=mstype.float32):
         stdv = 1 / math.sqrt(self.hidden_size)
         for weight in self.get_parameters():
-            weight.set_data(initializer(Uniform(stdv), weight.shape))
+            weight.set_data(initializer(Uniform(stdv), weight.shape, dtype))
 
 
 class RNNCell(RNNCellBase):
@@ -181,13 +185,14 @@ class RNNCell(RNNCellBase):
     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_ih` and `b_hh`. Default: True.
-        nonlinearity (str): The non-linearity to use. Can be either `tanh` or `relu`. Default: `tanh`.
+        has_bias (bool): Whether the cell has bias `b_ih` and `b_hh`. Default: ``True`` .
+        nonlinearity (str): The non-linearity to use. Can be either ``"tanh"`` or ``"relu"`` .
+            Default: ``"tanh"`` .
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Inputs:
         - **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 :math:`(batch\_size, hidden\_size)` .
@@ -201,9 +206,11 @@ class RNNCell(RNNCellBase):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> net = nn.RNNCell(10, 16)
-        >>> x = Tensor(np.ones([5, 3, 10]).astype(np.float32))
-        >>> hx = Tensor(np.ones([3, 16]).astype(np.float32))
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> net = ms.nn.RNNCell(10, 16)
+        >>> x = ms.Tensor(np.ones([5, 3, 10]).astype(np.float32))
+        >>> hx = ms.Tensor(np.ones([3, 16]).astype(np.float32))
         >>> output = []
         >>> for i in range(5):
         ...     hx = net(x[i], hx)
@@ -213,8 +220,9 @@ class RNNCell(RNNCellBase):
     """
     _non_linearity = ['tanh', 'relu']
 
-    def __init__(self, input_size: int, hidden_size: int, has_bias: bool = True, nonlinearity: str = "tanh"):
-        super().__init__(input_size, hidden_size, has_bias, num_chunks=1)
+    def __init__(self, input_size: int, hidden_size: int, has_bias: bool = True, nonlinearity: str = "tanh",
+                 dtype=mstype.float32):
+        super().__init__(input_size, hidden_size, has_bias, num_chunks=1, dtype=dtype)
         validator.check_value_type("nonlinearity", nonlinearity, [str], self.cls_name)
         validator.check_string(nonlinearity, self._non_linearity, "nonlinearity", self.cls_name)
         self.nonlinearity = nonlinearity
@@ -263,15 +271,16 @@ class LSTMCell(RNNCellBase):
     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_ih` and `b_hh`. Default: True.
+        has_bias (bool): Whether the cell has bias `b_ih` and `b_hh`. Default: ``True`` .
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Inputs:
-        - **x** (Tensor) - Tensor of shape :math:`(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 :math:`(batch\_size, hidden\_size)`. The data type of `hx` must be the same as `x`.
+          and shape :math:`(batch\_size, hidden\_size)` .
 
     Outputs:
-        - **hx'** (Tensor) - A tuple of two Tensors (h', c') both of data shape :math:`(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.
@@ -281,10 +290,12 @@ class LSTMCell(RNNCellBase):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> net = nn.LSTMCell(10, 16)
-        >>> x = Tensor(np.ones([5, 3, 10]).astype(np.float32))
-        >>> h = Tensor(np.ones([3, 16]).astype(np.float32))
-        >>> c = Tensor(np.ones([3, 16]).astype(np.float32))
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> net = ms.nn.LSTMCell(10, 16)
+        >>> x = ms.Tensor(np.ones([5, 3, 10]).astype(np.float32))
+        >>> h = ms.Tensor(np.ones([3, 16]).astype(np.float32))
+        >>> c = ms.Tensor(np.ones([3, 16]).astype(np.float32))
         >>> output = []
         >>> for i in range(5):
         ...     hx = net(x[i], (h, c))
@@ -293,8 +304,9 @@ class LSTMCell(RNNCellBase):
         (3, 16)
     """
     @_check_lstmcell_init
-    def __init__(self, input_size: int, hidden_size: int, has_bias: bool = True):
-        super().__init__(input_size, hidden_size, has_bias, num_chunks=4)
+    def __init__(self, input_size: int, hidden_size: int, has_bias: bool = True,
+                 dtype=mstype.float32):
+        super().__init__(input_size, hidden_size, has_bias, num_chunks=4, dtype=dtype)
         self.support_non_tensor_inputs = True
 
     def construct(self, x, hx):
@@ -343,15 +355,15 @@ class GRUCell(RNNCellBase):
     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.
+        has_bias (bool): Whether the cell has bias `b_in` and `b_hn`. Default: ``True`` .
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Inputs:
-        - **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`.
+        - **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)` .
 
     Outputs:
-        - **hx'** (Tensor) - Tensor of shape :math:`(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.
@@ -361,9 +373,11 @@ class GRUCell(RNNCellBase):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> net = nn.GRUCell(10, 16)
-        >>> x = Tensor(np.ones([5, 3, 10]).astype(np.float32))
-        >>> hx = Tensor(np.ones([3, 16]).astype(np.float32))
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> net = ms.nn.GRUCell(10, 16)
+        >>> x = ms.Tensor(np.ones([5, 3, 10]).astype(np.float32))
+        >>> hx = ms.Tensor(np.ones([3, 16]).astype(np.float32))
         >>> output = []
         >>> for i in range(5):
         ...     hx = net(x[i], hx)
@@ -371,8 +385,9 @@ class GRUCell(RNNCellBase):
         >>> print(output[0].shape)
         (3, 16)
     """
-    def __init__(self, input_size: int, hidden_size: int, has_bias: bool = True):
-        super().__init__(input_size, hidden_size, has_bias, num_chunks=3)
+    def __init__(self, input_size: int, hidden_size: int, has_bias: bool = True,
+                 dtype=mstype.float32):
+        super().__init__(input_size, hidden_size, has_bias, num_chunks=3, dtype=dtype)
 
     def construct(self, x, hx):
         _check_is_tensor('x', x, self.cls_name)

mindspore/nn/layer/rnns.py

@@ -35,10 +35,9 @@ from mindspore.nn.layer.rnn_cells import _rnn_relu_cell, _rnn_tanh_cell, _gru_ce
 __all__ = ['LSTM', 'GRU', 'RNN']
 
 
-@_primexpr
 def _init_state(shape, dtype, is_lstm):
-    hx = P.Zeros()(shape, dtype)
-    cx = P.Zeros()(shape, dtype)
+    hx = P.zeros(shape, dtype)
+    cx = P.zeros(shape, dtype)
     if is_lstm:
         return (hx, cx)
     return hx
@@ -58,8 +57,8 @@ def _check_input_dtype_same_and_valid(args_name, args_value, valid_values, cls_n
 @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):
-        raise TypeError(f"For '{cls_name}', the '{param_name}' must be '{mstype.tensor_type}', "
+    if input_data is not None and not isinstance(P.typeof(input_data), mstype.TensorType):
+        raise TypeError(f"For '{cls_name}', the '{param_name}' must be '{mstype.TensorType}', "
                         f"but got '{P.typeof(input_data)}'")
 
 
@@ -260,8 +259,8 @@ class _DynamicGRUAscend(Cell):
     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)
+            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), \
@@ -344,7 +343,7 @@ class _DynamicLSTMAscend(Cell):
         w_hh = self.concat_dim0((w_hh_i, w_hh_g, w_hh_f, w_hh_o))
         weight = self.concat_dim1((w_ih, w_hh))
         if b_ih is None:
-            bias = P.Zeros()(w_ih.shape[0], w_ih.dtype)
+            bias = P.zeros(w_ih.shape[0], w_ih.dtype)
         else:
             b_ih_i, b_ih_f, b_ih_g, b_ih_o = self.split(b_ih)
             b_hh_i, b_hh_f, b_hh_g, b_hh_o = self.split(b_hh)
@@ -373,7 +372,7 @@ class _RNNBase(Cell):
     '''Basic class for RNN operators'''
 
     def __init__(self, mode, input_size, hidden_size, num_layers=1, has_bias=True,
-                 batch_first=False, dropout=0., bidirectional=False):
+                 batch_first=False, dropout=0., bidirectional=False, dtype=mstype.float32):
         super().__init__()
         validator.check_positive_int(hidden_size, "hidden_size", self.cls_name)
         validator.check_positive_int(input_size, "input_size", self.cls_name)
@@ -436,17 +435,17 @@ class _RNNBase(Cell):
                 suffix = '_reverse' if direction == 1 else ''
 
                 self.w_ih_list.append(Parameter(
-                    Tensor(np.random.uniform(-stdv, stdv, (gate_size, layer_input_size)).astype(np.float32)),
-                    name='weight_ih_l{}{}'.format(layer, suffix)))
+                    Tensor(np.random.uniform(-stdv, stdv, (gate_size, layer_input_size)).astype(np.float32),
+                           dtype=dtype), name='weight_ih_l{}{}'.format(layer, suffix)))
                 self.w_hh_list.append(Parameter(
-                    Tensor(np.random.uniform(-stdv, stdv, (gate_size, hidden_size)).astype(np.float32)),
-                    name='weight_hh_l{}{}'.format(layer, suffix)))
+                    Tensor(np.random.uniform(-stdv, stdv, (gate_size, hidden_size)).astype(np.float32),
+                           dtype=dtype), name='weight_hh_l{}{}'.format(layer, suffix)))
                 if has_bias:
                     self.b_ih_list.append(Parameter(
-                        Tensor(np.random.uniform(-stdv, stdv, (gate_size)).astype(np.float32)),
+                        Tensor(np.random.uniform(-stdv, stdv, (gate_size)).astype(np.float32), dtype=dtype),
                         name='bias_ih_l{}{}'.format(layer, suffix)))
                     self.b_hh_list.append(Parameter(
-                        Tensor(np.random.uniform(-stdv, stdv, (gate_size)).astype(np.float32)),
+                        Tensor(np.random.uniform(-stdv, stdv, (gate_size)).astype(np.float32), dtype=dtype),
                         name='bias_hh_l{}{}'.format(layer, suffix)))
         self.w_ih_list = ParameterTuple(self.w_ih_list)
         self.w_hh_list = ParameterTuple(self.w_hh_list)
@@ -579,9 +578,7 @@ class _RNNBase(Cell):
 
 class RNN(_RNNBase):
     r"""
-    Stacked Elman RNN layers.
-
-    Apply RNN layer with :math:`\tanh` or :math:`\text{ReLU}` non-linearity to the input.
+    Stacked Elman RNN layers, applying RNN layer with :math:`\tanh` or :math:`\text{ReLU}` non-linearity to the input.
 
     For each element in the input sequence, each layer computes the following function:
 
@@ -596,23 +593,23 @@ class RNN(_RNNBase):
     Args:
         input_size (int): Number of features of input.
         hidden_size (int):  Number of features of hidden layer.
-        num_layers (int): Number of layers of stacked RNN. Default: 1.
+        num_layers (int): Number of layers of stacked RNN. Default: ``1`` .
         nonlinearity (str): The non-linearity to use. Can be either ``'tanh'`` or ``'relu'``. Default: ``'tanh'``
-        has_bias (bool): Whether the cell has bias `b_ih` and `b_hh`. Default: True.
-        batch_first (bool): Specifies whether the first dimension of input `x` is batch_size. Default: False.
+        has_bias (bool): Whether the cell has bias :math:`b_{ih}` and :math:`b_{hh}`. Default: ``True`` .
+        batch_first (bool): Specifies whether the first dimension of input `x` is batch_size. Default: ``False`` .
         dropout (float): If not 0.0, append `Dropout` layer on the outputs of each
-            RNN layer except the last layer. Default 0.0. The range of dropout is [0.0, 1.0).
+            RNN layer except the last layer. Default ``0.0`` . The range of dropout is [0.0, 1.0).
         bidirectional (bool): Specifies whether it is a bidirectional RNN,
-            num_directions=2 if bidirectional=True otherwise 1. Default: False.
+            num_directions=2 if bidirectional=True otherwise 1. Default: ``False`` .
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Inputs:
         - **x** (Tensor) - Tensor of data type mindspore.float32 or mindspore.float16 and
           shape :math:`(seq\_len, batch\_size, input\_size)` or :math:`(batch\_size, seq\_len, input\_size)` .
         - **hx** (Tensor) - Tensor of data type mindspore.float32 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:`(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.
@@ -635,9 +632,11 @@ class RNN(_RNNBase):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> net = nn.RNN(10, 16, 2, has_bias=True, batch_first=True, bidirectional=False)
-        >>> x = Tensor(np.ones([3, 5, 10]).astype(np.float32))
-        >>> h0 = Tensor(np.ones([1 * 2, 3, 16]).astype(np.float32))
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> net = ms.nn.RNN(10, 16, 2, has_bias=True, batch_first=True, bidirectional=False)
+        >>> x = ms.Tensor(np.ones([3, 5, 10]).astype(np.float32))
+        >>> h0 = ms.Tensor(np.ones([1 * 2, 3, 16]).astype(np.float32))
         >>> output, hn = net(x, h0)
         >>> print(output.shape)
         (3, 5, 16)
@@ -696,22 +695,22 @@ class GRU(_RNNBase):
     Args:
         input_size (int): Number of features of input.
         hidden_size (int):  Number of features of hidden layer.
-        num_layers (int): Number of layers of stacked GRU. Default: 1.
-        has_bias (bool): Whether the cell has bias `b_in` and `b_hn`. Default: True.
-        batch_first (bool): Specifies whether the first dimension of input `x` is batch_size. Default: False.
+        num_layers (int): Number of layers of stacked GRU. Default: ``1`` .
+        has_bias (bool): Whether the cell has bias `b_in` and `b_hn`. Default: ``True`` .
+        batch_first (bool): Specifies whether the first dimension of input `x` is batch_size. Default: ``False`` .
         dropout (float): If not 0.0, append `Dropout` layer on the outputs of each
-            GRU layer except the last layer. Default 0.0. The range of dropout is [0.0, 1.0).
+            GRU layer except the last layer. Default ``0.0`` . The range of dropout is [0.0, 1.0).
         bidirectional (bool): Specifies whether it is a bidirectional GRU,
-            num_directions=2 if bidirectional=True otherwise 1. Default: False.
+            num_directions=2 if bidirectional=True otherwise 1. Default: ``False`` .
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     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 (seq\_len, batch\_size, `input\_size`) or :math:`(batch\_size, seq\_len, input\_size)`.
         - **hx** (Tensor) - Tensor of data type mindspore.float32 or mindspore.float16 and
-          shape (num_directions * `num_layers`, batch_size, `hidden_size`). The data type of `hx` must be the same as
-          `x`.
+          shape :math:`(num\_directions * num\_layers, batch\_size, hidden\_size)`.
         - **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:`(\text{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.
@@ -719,9 +718,9 @@ class GRU(_RNNBase):
     Outputs:
         Tuple, a tuple contains (`output`, `h_n`).
 
-        - **output** (Tensor) - Tensor of shape (seq_len, batch_size, num_directions * `hidden_size`) or
-          (batch_size, seq_len, num_directions * `hidden_size`).
-        - **hx_n** (Tensor) - Tensor of shape (num_directions * `num_layers`, batch_size, `hidden_size`).
+        - **output** (Tensor) - Tensor of shape :math:`(seq\_len, batch\_size, num\_directions * hidden\_size)` or
+          :math:`(batch\_size, seq\_len, num\_directions * hidden\_size)`.
+        - **hx_n** (Tensor) - Tensor of shape :math:`(num\_directions * num\_layers, batch\_size, hidden\_size)`.
 
     Raises:
         TypeError: If `input_size`, `hidden_size` or `num_layers` is not an int.
@@ -733,9 +732,11 @@ class GRU(_RNNBase):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> net = nn.GRU(10, 16, 2, has_bias=True, batch_first=True, bidirectional=False)
-        >>> x = Tensor(np.ones([3, 5, 10]).astype(np.float32))
-        >>> h0 = Tensor(np.ones([1 * 2, 3, 16]).astype(np.float32))
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> net = ms.nn.GRU(10, 16, 2, has_bias=True, batch_first=True, bidirectional=False)
+        >>> x = ms.Tensor(np.ones([3, 5, 10]).astype(np.float32))
+        >>> h0 = ms.Tensor(np.ones([1 * 2, 3, 16]).astype(np.float32))
         >>> output, hn = net(x, h0)
         >>> print(output.shape)
         (3, 5, 16)
@@ -794,22 +795,22 @@ class LSTM(_RNNBase):
     Args:
         input_size (int): Number of features of input.
         hidden_size (int):  Number of features of hidden layer.
-        num_layers (int): Number of layers of stacked LSTM . Default: 1.
-        has_bias (bool): Whether the cell has bias `b_ih` and `b_hh`. Default: True.
-        batch_first (bool): Specifies whether the first dimension of input `x` is batch_size. Default: False.
+        num_layers (int): Number of layers of stacked LSTM . Default: ``1`` .
+        has_bias (bool): Whether the cell has bias :math:`b_{ih}` and :math:`b_{hh}`. Default: ``True`` .
+        batch_first (bool): Specifies whether the first dimension of input `x` is batch_size. Default: ``False`` .
         dropout (float, int): If not 0, append `Dropout` layer on the outputs of each
-            LSTM layer except the last layer. Default 0. The range of dropout is [0.0, 1.0).
+            LSTM layer except the last layer. Default ``0`` . The range of dropout is [0.0, 1.0).
         bidirectional (bool): Specifies whether it is a bidirectional LSTM,
-            num_directions=2 if bidirectional=True otherwise 1. Default: False.
+            num_directions=2 if bidirectional=True otherwise 1. Default: ``False`` .
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Inputs:
         - **x** (Tensor) - Tensor of data type mindspore.float32 or mindspore.float16 and
-          shape :math:`(seq\_len, batch\_size, input\_size)` or :math:`(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 :math:`(num\_directions * num\_layers, batch\_size, hidden\_size)`.
-          The data type of `hx` must be the same as `x`.
+          or mindspore.float16 and shape :math:`(num\_directions * num\_layers, batch\_size, hidden\_size)` .
         - **seq_length** (Tensor) - The length of each sequence in an input batch.
-          Tensor of shape :math:`(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.
@@ -831,10 +832,12 @@ class LSTM(_RNNBase):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> net = nn.LSTM(10, 16, 2, has_bias=True, batch_first=True, bidirectional=False)
-        >>> x = Tensor(np.ones([3, 5, 10]).astype(np.float32))
-        >>> h0 = Tensor(np.ones([1 * 2, 3, 16]).astype(np.float32))
-        >>> c0 = Tensor(np.ones([1 * 2, 3, 16]).astype(np.float32))
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> net = ms.nn.LSTM(10, 16, 2, has_bias=True, batch_first=True, bidirectional=False)
+        >>> x = ms.Tensor(np.ones([3, 5, 10]).astype(np.float32))
+        >>> h0 = ms.Tensor(np.ones([1 * 2, 3, 16]).astype(np.float32))
+        >>> c0 = ms.Tensor(np.ones([1 * 2, 3, 16]).astype(np.float32))
         >>> output, (hn, cn) = net(x, (h0, c0))
         >>> print(output.shape)
         (3, 5, 16)