mindspore 2.3.0__cp310-cp310-win_amd64.whl → 2.4.1__cp310-cp310-win_amd64.whl

mindspore/nn/layer/activation.py

@@ -33,6 +33,7 @@ __all__ = ['Softmin',
            'Softmax',
            'Softmax2d',
            'LogSoftmax',
+           'LogSoftmaxExt',
            'ReLU',
            'ReLU6',
            'RReLU',
@@ -46,6 +47,7 @@ __all__ = ['Softmin',
            'Sigmoid',
            'Softsign',
            'PReLU',
+           'PReLUExt',
            'get_activation',
            'LeakyReLU',
            'HSigmoid',
@@ -279,6 +281,35 @@ class Softmax(Cell):
         return self.softmax(input)
 
 
+class SoftmaxExt(Cell):
+    r"""
+    Applies the Softmax function to an n-dimensional input Tensor.
+
+    For details, please refer to :func:`mindspore.mint.nn.functional.softmax`.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, nn
+        >>> import numpy as np
+        >>> input = Tensor(np.array([-1, -2, 0, 2, 1]), mindspore.float16)
+        >>> softmax = nn.SoftmaxExt()
+        >>> output = softmax(input)
+        >>> print(output)
+        [0.03168 0.01166 0.0861  0.636   0.2341 ]
+    """
+
+    def __init__(self, dim=None):
+        """Initialize Softmax."""
+        super(SoftmaxExt, self).__init__()
+        self.dim = dim
+
+    def construct(self, input):
+        return ops.function.nn_func.softmax_ext(input, self.dim)
+
+
 class LogSoftmax(Cell):
     r"""
     Applies the LogSoftmax function to n-dimensional input tensor element-wise.
@@ -329,6 +360,51 @@ class LogSoftmax(Cell):
         return self.log_softmax(x)
 
 
+class LogSoftmaxExt(Cell):
+    r"""
+    Applies the Log Softmax function to the input tensor on the specified axis.
+    Supposes a slice in the given axis, :math:`x` for each element :math:`x_i`,
+    the Log Softmax function is shown as follows:
+
+    .. math::
+        \text{output}(x_i) = \log \left(\frac{\exp(x_i)} {\sum_{j = 0}^{N-1}\exp(x_j)}\right),
+
+    where :math:`N` is the length of the Tensor.
+
+    Args:
+        dim (int, optional): The axis to perform the Log softmax operation. Default: ``None`` .
+
+    Returns:
+        Tensor, with the same shape as the input.
+
+    Raises:
+        ValueError: If `dim` is not in range [-len(input.shape), len(input.shape)).
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, nn
+        >>> import numpy as np
+        >>> x = Tensor(np.array([[-1.0, 4.0, -8.0], [2.0, -5.0, 9.0]]), mindspore.float32)
+        >>> log_softmax = nn.LogSoftmaxExt(dim=-1)
+        >>> output = log_softmax(x)
+        >>> print(output)
+        [[-5.00672150e+00 -6.72150636e-03 -1.20067215e+01]
+         [-7.00091219e+00 -1.40009127e+01 -9.12250078e-04]]
+    """
+
+    def __init__(self, dim=None):
+        """Initialize LogSoftmaxExt."""
+        super(LogSoftmaxExt, self).__init__()
+        self.log_softmax = P.LogSoftmaxExt()
+        self.dim = dim
+
+    def construct(self, x):
+        return self.log_softmax(x, dim=self.dim)
+
+
 class ELU(Cell):
     r"""
     Applies the exponential linear unit function element-wise.
@@ -434,8 +510,8 @@ class ReLU(Cell):
         super(ReLU, self).__init__()
         self.relu = P.ReLU()
 
-    def construct(self, x):
-        return self.relu(x)
+    def construct(self, input):
+        return self.relu(input)
 
 
 class ReLU6(Cell):
@@ -898,6 +974,13 @@ class GELU(Cell):
     Outputs:
         Tensor, with the same type and shape as the `x`.
 
+    Note:
+        when calculating the input gradient of GELU with an input value of infinity, there are differences
+        in the output of the backward between ``Ascend`` and ``GPU``.
+        when x is -inf, the computation result of ``Ascend`` is 0, and the computation result of ``GPU`` is Nan.
+        when x is inf, the computation result of ``Ascend`` is dy, and the computation result of ``GPU`` is Nan.
+        In mathematical terms, the result of Ascend has higher precision.
+
     Raises:
         TypeError: If dtype of `x` is not one of float16, float32, or float64.
 
@@ -1164,14 +1247,85 @@ class PReLU(Cell):
         return self.prelu(x, F.cast(self.w, x.dtype))
 
 
+class PReLUExt(Cell):
+    r"""
+    Applies PReLU activation function element-wise.
+
+    PReLU is defined as:
+
+    .. math::
+
+        PReLU(x_i)= \max(0, x_i) + w * \min(0, x_i),
+
+    where :math:`x_i` is an element of an channel of the input.
+
+    Here :math:`w` is a learnable parameter with a default initial value 0.25.
+    Parameter :math:`w` has dimensionality of the argument channel. If called without argument
+    channel, a single parameter :math:`w` will be shared across all channels.
+
+    PReLU Activation Function Graph:
+
+    .. image:: ../images/PReLU2.png
+        :align: center
+
+    .. note::
+        Channel dim is the 2nd dim of input. When input has dims < 2, then there is
+        no channel dim and the number of channels = 1.
+
+    Args:
+        num_parameters (int): number of `w` to learn. Although it takes an int as input,
+            there is only two legitimate values: 1, or the number of channels at Tensor `input`. Default: ``1`` .
+        init (float): the initial value of `w`. Default: ``0.25`` .
+        dtype (mindspore.dtype, optional): the type of `w`. Default: ``None`` . Supported data type
+            is {float16, float32, bfloat16}.
+
+    Inputs:
+        - **input** (Tensor) - The input of PReLU.
+
+    Outputs:
+        Tensor, with the same dtype and shape as the `input`.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, nn
+        >>> import numpy as np
+        >>> x = Tensor(np.array([[[[0.1, 0.6], [0.9, 0.9]]]]), mindspore.float32)
+        >>> prelu = nn.PReLUExt()
+        >>> output = prelu(x)
+        >>> print(output)
+        [[[[0.1 0.6]
+           [0.9 0.9]]]]
+
+    """
+
+    def __init__(self, num_parameters=1, init=0.25, dtype=None):
+        """Initialize PReLUExt."""
+        super(PReLUExt, self).__init__()
+        tmp = np.empty((num_parameters,), dtype=np.float32)
+        tmp.fill(init)
+        w = Tensor(tmp, dtype=dtype)
+        self.weight = Parameter(w, name='weight')
+
+    def construct(self, input):
+        return ops.prelu(input, self.weight)
+
+
 class HSwish(Cell):
     r"""
-    Applies hswish-type activation element-wise.
+    Applies Hard Swish activation function element-wise.
 
     Hard swish is defined as:
 
     .. math::
-        \text{hswish}(x_{i}) = x_{i} * \frac{ReLU6(x_{i} + 3)}{6},
+        \text{Hardswish}(input) =
+        \begin{cases}
+        0, & \text{ if } input \leq -3, \\
+        input, & \text{ if } input \geq +3, \\
+        input*(input + 3)/6, & \text{ otherwise }
+        \end{cases}
 
     HSwish Activation Function Graph:
 
@@ -1179,14 +1333,14 @@ class HSwish(Cell):
         :align: center
 
     Inputs:
-        - **x** (Tensor) - The input of HSwish, data type must be float16 or float32.
-          The shape is :math:`(N,*)` where :math:`*` means, any number of additional dimensions.
+        - **input** (Tensor) - The input of HSwish.
 
     Outputs:
-        Tensor, with the same type and shape as the `x`.
+        Tensor, with the same type and shape as the `input`.
 
     Raises:
-        TypeError: If dtype of `x` is neither float16 nor float32.
+        TypeError: If `input` is not a tensor.
+        TypeError: If `input` is neither int nor float.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -1195,9 +1349,9 @@ class HSwish(Cell):
         >>> import mindspore
         >>> from mindspore import Tensor, nn
         >>> import numpy as np
-        >>> x = Tensor(np.array([-1, -2, 0, 2, 1]), mindspore.float16)
+        >>> input = Tensor(np.array([-1, -2, 0, 2, 1]), mindspore.float16)
         >>> hswish = nn.HSwish()
-        >>> result = hswish(x)
+        >>> result = hswish(input)
         >>> print(result)
         [-0.3333 -0.3333  0.      1.667   0.6665]
     """
@@ -1207,18 +1361,23 @@ class HSwish(Cell):
         super(HSwish, self).__init__()
         self.hswish = P.HSwish()
 
-    def construct(self, x):
-        return self.hswish(x)
+    def construct(self, input):
+        return self.hswish(input)
 
 
 class HSigmoid(Cell):
     r"""
-    Applies Hard sigmoid activation function element-wise.
+    Applies Hard Sigmoid activation function element-wise.
 
-    Hard sigmoid is defined as:
+    Hard Sigmoid is defined as:
 
     .. math::
-        \text{hsigmoid}(x_{i}) = \max(0, \min(1, \frac{x_{i} + 3}{6})),
+        \text{Hardsigmoid}(input) =
+        \begin{cases}
+        0, & \text{ if } input \leq -3, \\
+        1, & \text{ if } input \geq +3, \\
+        input/6 + 1/2, & \text{ otherwise }
+        \end{cases}
 
     HSigmoid Activation Function Graph:
 
@@ -1226,13 +1385,14 @@ class HSigmoid(Cell):
         :align: center
 
     Inputs:
-        - **input_x** (Tensor) - The input of HSigmoid. Tensor of any dimension.
+        - **input** (Tensor) - The input of HSigmoid.
 
     Outputs:
-        Tensor, with the same type and shape as the `input_x`.
+        Tensor, with the same type and shape as the `input`.
 
     Raises:
-        TypeError: If `input_x` is not a Tensor.
+        TypeError: If `input` is not a Tensor.
+        TypeError: If `input` is neither int nor float.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -1241,9 +1401,9 @@ class HSigmoid(Cell):
         >>> import mindspore
         >>> from mindspore import Tensor, nn
         >>> import numpy as np
-        >>> x = Tensor(np.array([-1, -2, 0, 2, 1]), mindspore.float16)
+        >>> input = Tensor(np.array([-1, -2, 0, 2, 1]), mindspore.float16)
         >>> hsigmoid = nn.HSigmoid()
-        >>> result = hsigmoid(x)
+        >>> result = hsigmoid(input)
         >>> print(result)
         [0.3333 0.1666 0.5    0.8335 0.6665]
     """
@@ -1253,8 +1413,8 @@ class HSigmoid(Cell):
         super(HSigmoid, self).__init__()
         self.hsigmoid = P.HSigmoid()
 
-    def construct(self, input_x):
-        return self.hsigmoid(input_x)
+    def construct(self, input):
+        return self.hsigmoid(input)
 
 
 class LogSigmoid(Cell):
@@ -1370,21 +1530,22 @@ class SoftShrink(Cell):
         :align: center
 
     Args:
-        lambd (float): the :math:`\lambda` must be no less than zero for the SoftShrink formulation.
-            Default: ``0.5`` .
+        lambd (number, optional): The threshold :math:`\lambda` defined by the Soft Shrink formula.
+            It should be greater than or equal to 0, default: ``0.5`` .
 
     Inputs:
-        - **input_x** (Tensor) - The input of SoftShrink with data type of float16 or float32.
-          Any number of additional dimensions.
+        - **input** (Tensor) - The input of Soft Shrink. Supported dtypes:
+
+          - Ascend: float16, float32, bfloat16.
+          - CPU/GPU: float16, float32.
 
     Outputs:
-        Tensor, has the same shape and data type as `input_x`.
+        Tensor, the same shape and data type as the input.
 
     Raises:
-        TypeError: If lambd is not a float.
-        TypeError: If input_x is not a Tensor.
-        TypeError: If dtype of input_x is neither float16 nor float32.
-        ValueError: If lambd is less than 0.
+        TypeError: If `lambd` is not a float, int or bool.
+        TypeError: If `input` is not a tensor.
+        TypeError: If dtype of `input` is not float16, float32 or bfloat16.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -1393,9 +1554,9 @@ class SoftShrink(Cell):
         >>> import mindspore
         >>> from mindspore import Tensor, nn
         >>> import numpy as np
-        >>> input_x = Tensor(np.array([[ 0.5297,  0.7871,  1.1754], [ 0.7836,  0.6218, -1.1542]]), mindspore.float16)
+        >>> input = Tensor(np.array([[ 0.5297,  0.7871,  1.1754], [ 0.7836,  0.6218, -1.1542]]), mindspore.float16)
         >>> softshrink = nn.SoftShrink()
-        >>> output = softshrink(input_x)
+        >>> output = softshrink(input)
         >>> print(output)
         [[ 0.02979  0.287    0.676  ]
          [ 0.2837   0.1216  -0.6543 ]]
@@ -1405,8 +1566,8 @@ class SoftShrink(Cell):
         super(SoftShrink, self).__init__()
         self.softshrink = P.SoftShrink(lambd)
 
-    def construct(self, input_x):
-        output = self.softshrink(input_x)
+    def construct(self, input):
+        output = self.softshrink(input)
         return output
 
 
@@ -1430,17 +1591,21 @@ class HShrink(Cell):
         :align: center
 
     Args:
-        lambd (float): The threshold :math:`\lambda` defined by the Hard Shrink formula. Default: ``0.5`` .
+        lambd (number, optional): The threshold :math:`\lambda` defined by the Hard Shrink formula. Default: ``0.5`` .
 
     Inputs:
-        - **input_x** (Tensor) - The input of Hard Shrink with data type of float16 or float32.
+        - **input** (Tensor) - The input of Hard Shrink. Supported dtypes:
+
+          - Ascend: float16, float32, bfloat16.
+          - CPU/GPU: float16, float32.
 
     Outputs:
         Tensor, the same shape and data type as the input.
 
     Raises:
-        TypeError: If `lambd` is not a float.
-        TypeError: If dtype of `input_x` is neither float16 nor float32.
+        TypeError: If `lambd` is not a float, int or bool.
+        TypeError: If `input` is not a tensor.
+        TypeError: If dtype of `input` is not float16, float32 or bfloat16.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -1449,20 +1614,20 @@ class HShrink(Cell):
         >>> import mindspore
         >>> from mindspore import Tensor, nn
         >>> import numpy as np
-        >>> input_x = Tensor(np.array([[ 0.5,  1,  2.0], [0.0533,0.0776,-2.1233]]), mindspore.float32)
+        >>> input = Tensor(np.array([[0.5, 1, 2.0], [0.0533, 0.0776, -2.1233]]), mindspore.float32)
         >>> hshrink = nn.HShrink()
-        >>> output = hshrink(input_x)
+        >>> output = hshrink(input)
         >>> print(output)
         [[ 0.      1.      2.    ]
-        [ 0.      0.     -2.1233]]
+         [ 0.      0.     -2.1233]]
     """
 
     def __init__(self, lambd=0.5):
         super(HShrink, self).__init__()
         self.hshrink = P.HShrink(lambd)
 
-    def construct(self, input_x):
-        return self.hshrink(input_x)
+    def construct(self, input):
+        return self.hshrink(input)
 
 
 class Threshold(Cell):
@@ -1602,6 +1767,7 @@ _activation = {
     'softmax': Softmax,
     'softmax2d': Softmax2d,
     'logsoftmax': LogSoftmax,
+    'logsoftmaxExt': LogSoftmaxExt,
     'relu': ReLU,
     'relu6': ReLU6,
     'rrelu': RReLU,
@@ -1615,6 +1781,7 @@ _activation = {
     'sigmoid': Sigmoid,
     'softsign': Softsign,
     'prelu': PReLU,
+    'preluExt': PReLUExt,
     'leakyrelu': LeakyReLU,
     'hswish': HSwish,
     'hsigmoid': HSigmoid,

mindspore/nn/layer/basic.py

@@ -40,7 +40,7 @@ from mindspore.common._decorator import deprecated
 from mindspore.ops.auto_generate import dropout_ext_op, fold_ext
 from mindspore.common.generator import default_generator
 
-__all__ = ['Dropout', 'Flatten', 'Dense', 'ClipByNorm', 'Norm', 'OneHot', 'Pad', 'Unfold', 'Tril', 'Triu',
+__all__ = ['Dropout', 'Flatten', 'Dense', 'Linear', 'ClipByNorm', 'Norm', 'OneHot', 'Pad', 'Unfold', 'Tril', 'Triu',
            'MatrixDiag', 'MatrixDiagPart', 'MatrixSetDiag', 'L1Regularizer', 'Dropout1d',
            'Dropout2d', 'Dropout3d', 'Upsample', 'Roll', 'Identity', 'Unflatten', 'DropoutExt']
 
@@ -510,8 +510,8 @@ class UpsampleExt(Cell):
         self.align_corners = align_corners
         self.recompute_scale_factor = recompute_scale_factor
 
-    def construct(self, x):
-        out = interpolate_ext(x, self.size, self.scale_factor, self.mode,
+    def construct(self, input):
+        out = interpolate_ext(input, self.size, self.scale_factor, self.mode,
                               self.align_corners, self.recompute_scale_factor)
         return out
 
@@ -579,11 +579,15 @@ class Identity(Cell):
     r"""
     A placeholder identity operator that returns the same as input.
 
+    Args:
+        args (Any): Any argument.
+        kwargs (Any): Any keyword argument.
+
     Inputs:
-        - **x** (Any) - The input of Identity.
+        - **input** (Any) - The input of Identity.
 
     Outputs:
-        The same as `x`.
+        The same as `input`.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -592,19 +596,19 @@ class Identity(Cell):
         >>> import mindspore
         >>> from mindspore import Tensor, nn
         >>> import numpy as np
-        >>> x = Tensor(np.array([1, 2, 3, 4]), mindspore.int64)
+        >>> input = Tensor(np.array([1, 2, 3, 4]), mindspore.int64)
         >>> net = nn.Identity()
-        >>> output = net(x)
+        >>> output = net(input)
         >>> print(output)
         [1 2 3 4]
     """
 
-    def __init__(self):
+    def __init__(self, *args, **kwargs):
         """Initialize Identity."""
         super(Identity, self).__init__()
 
-    def construct(self, x):
-        return x
+    def construct(self, input):
+        return input
 
 
 class Dense(Cell):
@@ -621,6 +625,9 @@ class Dense(Cell):
     data type as the :math:`X` created by the layer, and :math:`\text{bias}` is a bias vector
     with the same data type as the :math:`X` created by the layer (only if has_bias is True).
 
+    .. warning::
+        In PYNATIVE mode, if `bias` is ``False`` , the `x` cannot be greater than 6D.
+
     Args:
         in_channels (int): The number of channels in the input space.
         out_channels (int): The number of channels in the output space.
@@ -635,6 +642,8 @@ class Dense(Cell):
             layer. Both activation name, e.g. 'relu', and mindspore activation function, e.g. mindspore.ops.ReLU(),
             are supported. Default: ``None`` .
         dtype (:class:`mindspore.dtype`): Data type of Parameter. Default: ``mstype.float32`` .
+            When `weight_init` is Tensor, Parameter has the same data type as `weight_init` ,
+            in other cases, Parameter has the same data type as `dtype`, the same goes for `bias_init`.
 
     Inputs:
         - **x** (Tensor) - Tensor of shape :math:`(*, in\_channels)`. The `in_channels` in `Args` should be equal
@@ -651,6 +660,7 @@ class Dense(Cell):
                     is not equal to `out_channels` or shape[1] of `weight_init` is not equal to `in_channels`.
         ValueError: If length of shape of `bias_init` is not equal to 1
                     or shape[0] of `bias_init` is not equal to `out_channels`.
+        RuntimeError: If `bias` is ``False`` and `x` is greater than 6D in PYNATIVE mode.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -743,6 +753,123 @@ class Dense(Cell):
         return s
 
 
+class Linear(Cell):
+    r"""
+    The linear connected layer.
+
+    Applies linear connected layer for the input. This layer implements the operation as:
+
+    .. math::
+        \text{outputs} = X * kernel + bias
+
+    .. warning::
+        In PYNATIVE mode, if `bias` is ``False`` , the `x` cannot be greater than 6D.
+
+    where :math:`X` is the input tensors, :math:`\text{kernel}` is a weight matrix with the same
+    data type as the :math:`X` created by the layer, and :math:`\text{bias}` is a bias vector
+    with the same data type as the :math:`X` created by the layer (only if has_bias is True).
+
+    Args:
+        in_features (int): The number of features in the input space.
+        out_features (int): The number of features in the output space.
+        bias (bool): Specifies whether the layer uses a bias vector :math:`\text{bias}`. Default: ``True``.
+        weight_init (Union[Tensor, str, Initializer, numbers.Number]): The trainable weight_init parameter. The dtype
+            is same as `x`. The values of str refer to the function `initializer`. Default: ``None`` ,
+            weight will be initialized using HeUniform.
+        bias_init (Union[Tensor, str, Initializer, numbers.Number]): The trainable bias_init parameter. The dtype is
+            same as `x`. The values of str refer to the function `initializer`. Default: ``None`` ,
+            bias will be initialized using Uniform.
+        dtype (:class:`mindspore.dtype`): Data type of Parameter. Default: ``None`` .
+            If `dtype` is ``None`` , `dtype` is set to ``mstype.float32`` when initializing the method.
+            When `weight_init` is Tensor, Parameter has the same data type as `weight_init` ,
+            in other cases, Parameter has the same data type as `dtype`, the same goes for `bias_init`.
+
+    Inputs:
+        - **x** (Tensor) - Tensor of shape :math:`(*, in\_features)`. The `in_features` in `Args` should be equal
+          to :math:`in\_features` in `Inputs`.
+
+    Outputs:
+        Tensor of shape :math:`(*, out\_features)`.
+
+    Raises:
+        TypeError: If `in_features` or `out_features` is not an int.
+        TypeError: If `bias` is not a bool.
+        ValueError: If length of shape of `weight_init` is not equal to 2 or shape[0] of `weight_init`
+                    is not equal to `out_features` or shape[1] of `weight_init` is not equal to `in_features`.
+        ValueError: If length of shape of `bias_init` is not equal to 1
+                    or shape[0] of `bias_init` is not equal to `out_features`.
+        RuntimeError: If `bias` is ``False`` and `x` is greater than 6D in PYNATIVE mode.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor
+        >>> from mindspore import nn
+        >>> import numpy as np
+        >>> x = Tensor(np.array([[180, 234, 154], [244, 48, 247]]), mindspore.float32)
+        >>> net = nn.mint.nn.Linear(3, 4)
+        >>> output = net(x)
+        >>> print(output.shape)
+        (2, 4)
+    """
+
+    @cell_attr_register(attrs=['has_bias'])
+    def __init__(self,
+                 in_features,
+                 out_features,
+                 bias=True,
+                 weight_init=None,
+                 bias_init=None,
+                 dtype=None):
+        """Initialize Linear."""
+        super(Linear, self).__init__()
+        self.in_features = Validator.check_positive_int(
+            in_features, "in_features", self.cls_name)
+        self.out_features = Validator.check_positive_int(
+            out_features, "out_features", self.cls_name)
+        self.has_bias = Validator.check_bool(
+            bias, "has_bias", self.cls_name)
+        self.dense = P.Dense()
+        if dtype is None:
+            dtype = mstype.float32
+        if isinstance(weight_init, Tensor):
+            if weight_init.ndim != 2 or weight_init.shape[0] != out_features or \
+                    weight_init.shape[1] != in_features:
+                raise ValueError(f"For '{self.cls_name}', weight init shape error. The ndim of 'weight_init' must "
+                                 f"be equal to 2, and the first dim must be equal to 'out_features', and the "
+                                 f"second dim must be equal to 'in_features'. But got 'weight_init': {weight_init}, "
+                                 f"'out_features': {out_features}, 'in_features': {in_features}.")
+        if weight_init is None:
+            weight_init = HeUniform(math.sqrt(5))
+        self.weight = Parameter(initializer(
+            weight_init, [out_features, in_features], dtype=dtype), name="weight")
+
+        self.bias = None
+        if self.has_bias:
+            if isinstance(bias_init, Tensor):
+                if bias_init.ndim != 1 or bias_init.shape[0] != out_features:
+                    raise ValueError(f"For '{self.cls_name}', bias init shape error. The ndim of 'bias_init' must "
+                                     f"be equal to 1, and the first dim must be equal to 'out_features'. But got "
+                                     f"'bias_init': {bias_init}, 'out_features': {out_features}.")
+            if bias_init is None:
+                bound = 1 / math.sqrt(in_features)
+                bias_init = Uniform(scale=bound)
+            self.bias = Parameter(initializer(
+                bias_init, [out_features], dtype=dtype), name="bias")
+
+    def construct(self, x):
+        x = self.dense(x, self.weight, self.bias)
+        return x
+
+    def extend_repr(self):
+        s = f'input_features={self.in_features}, output_features={self.out_features}'
+        if self.has_bias:
+            s += f', has_bias={self.has_bias}'
+        return s
+
+
 @constexpr
 def _is_equal_one(x):
     if x is None: