mindspore 2.3.0rc1__cp37-cp37m-manylinux1_x86_64.whl → 2.3.0rc2__cp37-cp37m-manylinux1_x86_64.whl

mindspore/nn/cell.py

@@ -20,10 +20,9 @@ import inspect
 import os
 import time
 from collections import OrderedDict
-from types import FunctionType, MethodType
 import numpy
 
-from mindspore._checkparam import args_type_check
+from mindspore._checkparam import args_type_check, check_hook_fn
 from mindspore.common._auto_dynamic import is_auto_dynamic, convert_inputs_to_dynamic
 from mindspore import log as logger
 from mindspore.common.parameter import PARAMETER_NAME_DEFAULT
@@ -34,7 +33,7 @@ from mindspore._c_expression import init_pipeline, update_func_graph_hyper_param
 from mindspore import _checkparam as Validator
 from mindspore.common import dtype as mstype
 from mindspore.common.api import _cell_graph_executor, _pynative_executor, _get_args_for_run, cells_compile_cache
-from mindspore.common.api import _generate_branch_control_input
+from mindspore.common.api import _generate_branch_control_input, _convert_python_data, _get_args_for_run_predict
 from mindspore.common.parameter import Parameter, ParameterTuple
 from mindspore.common.tensor import Tensor
 from mindspore.ops.operations import Cast
@@ -43,6 +42,7 @@ from mindspore.ops.operations import _inner_ops as inner
 from mindspore.parallel.shard import Shard
 from mindspore._check_jit_forbidden_api import jit_forbidden_register
 from mindspore.common._decorator import deprecated
+from mindspore.common._register_for_recompute import recompute_registry
 
 
 class Cell(Cell_):
@@ -125,11 +125,13 @@ class Cell(Cell_):
         self._create_time = int(time.time() * 1e9)
         self.arguments_key = ""
         self.compile_cache = set()
+        self.phase_cache = dict()
         cells_compile_cache[id(self)] = self.compile_cache
         self.parameter_broadcast_done = False
         self._id = 1
         self.exist_names = set("")
         self.exist_objs = set()
+        self.recompute_cell = None
         init_pipeline()
 
         # call gc to release GE session resources used by non-used cell objects
@@ -217,7 +219,7 @@ class Cell(Cell_):
 
         Tutorial Examples:
             - `Cell and Parameter - Custom Cell Reverse
-              <https://mindspore.cn/tutorials/en/r2.3.q1/advanced/modules/layer.html#custom-cell-reverse>`_
+              <https://mindspore.cn/tutorials/en/master/advanced/modules/layer.html#custom-cell-reverse>`_
         """
         return self._bprop_debug
 
@@ -415,7 +417,7 @@ class Cell(Cell_):
             elif isinstance(item, float):
                 res.append(self.cast(item, dst_type))
             elif hasattr(item, "dtype") and item.dtype in \
-                {mstype.float16, mstype.float32, mstype.float64, mstype.bfloat16} and item.dtype != dst_type:
+                    {mstype.float16, mstype.float32, mstype.float64, mstype.bfloat16} and item.dtype != dst_type:
                 res.append(self.cast(item, dst_type))
             else:
                 res.append(item)
@@ -474,7 +476,10 @@ class Cell(Cell_):
         elif hasattr(self, "_shard_fn"):
             output = self._shard_fn(*cast_inputs, **kwargs)
         else:
-            output = self.construct(*cast_inputs, **kwargs)
+            if self.recompute_cell is not None:
+                output = self.recompute_cell(*cast_inputs, **kwargs)
+            else:
+                output = self.construct(*cast_inputs, **kwargs)
         if self._enable_forward_hook:
             output = self._run_forward_hook(cast_inputs, output)
         return output
@@ -659,6 +664,16 @@ class Cell(Cell_):
             self.check_names_and_refresh_name()
             self._is_check_and_refresh = True
 
+    def _predict(self, *args, **kwargs):
+        if not hasattr(self, "phase"):
+            return False, None
+        if (self.phase == "prefill" or self.phase == 'increment') and self.phase in self.phase_cache:
+            new_args = _get_args_for_run_predict(self, args, kwargs, self._compile_args)
+            res = _cell_graph_executor._graph_executor(tuple(new_args), self.phase_cache[self.phase])
+            res = _convert_python_data(res)
+            return True, res
+        return False, None
+
     def __call__(self, *args, **kwargs):
         # Run in Graph mode.
         if os.getenv("MS_JIT") != '0' and context._get_mode() == context.GRAPH_MODE:
@@ -667,7 +682,12 @@ class Cell(Cell_):
                 bound_arguments.apply_defaults()
                 args = bound_arguments.args
                 kwargs = bound_arguments.kwargs
+
+            predict_compiled, res = self._predict(*args, **kwargs)
+            if predict_compiled:
+                return res
             self._check_construct_args(*args)
+
             if self._hook_fn_registered():
                 logger.warning(f"For 'Cell', it's not support hook function in graph mode. If you want to use hook "
                                f"function, please use context.set_context to set pynative mode.")
@@ -964,7 +984,6 @@ class Cell(Cell_):
             return self._dynamic_shape_inputs
         return args
 
-
     def compile(self, *args, **kwargs):
         """
         Compile Cell as a computation graph, the input must be consistent with the input defined in construct.
@@ -1335,7 +1354,7 @@ class Cell(Cell_):
 
         Tutorial Examples:
             - `Model Training - Optimizer
-              <https://mindspore.cn/tutorials/en/r2.3.q1/beginner/train.html#optimizer>`_
+              <https://mindspore.cn/tutorials/en/master/beginner/train.html#optimizer>`_
         """
         return list(filter(lambda x: x.requires_grad, self.get_parameters(expand=recurse)))
 
@@ -1446,7 +1465,7 @@ class Cell(Cell_):
 
         Tutorial Examples:
             - `Building a Network - Model Parameters
-              <https://mindspore.cn/tutorials/en/r2.3.q1/beginner/model.html#model-parameters>`_
+              <https://mindspore.cn/tutorials/en/master/beginner/model.html#model-parameters>`_
         """
         cells = []
         if expand:
@@ -1785,7 +1804,7 @@ class Cell(Cell_):
         accelerate the algorithm in the algorithm library.
 
         If `boost_type` is not in the algorithm library, please view the algorithm in the algorithm library through
-        `algorithm library <https://gitee.com/mindspore/mindspore/tree/r2.3.q1/mindspore/python/mindspore/boost>`_.
+        `algorithm library <https://gitee.com/mindspore/mindspore/tree/master/mindspore/python/mindspore/boost>`_.
 
         Note:
             Some acceleration algorithms may affect the accuracy of the network, please choose carefully.
@@ -1842,7 +1861,7 @@ class Cell(Cell_):
 
         Tutorial Examples:
             - `Model Training - Implementing Training and Evaluation
-              <https://mindspore.cn/tutorials/en/r2.3.q1/beginner/train.html#training-and-evaluation>`_
+              <https://mindspore.cn/tutorials/en/master/beginner/train.html#training-and-evaluation>`_
         """
         if mode:
             self._phase = 'train'
@@ -1936,8 +1955,8 @@ class Cell(Cell_):
             hook_fn (function): Python function. Forward pre hook function.
 
         Returns:
-            Handle, it is an instance of `mindspore.common.hook_handle.HookHandle` and corresponding to the `hook_fn` .
-            The handle can be used to remove the added `hook_fn` by calling `handle.remove()` .
+            A handle corresponding to the `hook_fn` . The handle can be used to remove the added `hook_fn` by calling
+            `handle.remove()` .
 
         Raises:
             TypeError: If the `hook_fn` is not a function of python.
@@ -1972,17 +1991,8 @@ class Cell(Cell_):
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32,
             value= [ 2.00000000e+00]))
         """
-        if context.get_context("mode") != context.PYNATIVE_MODE:
-            logger.warning(f"'register_forward_pre_hook' function is only supported in pynative mode, you can use "
-                           f"context.set_context to set pynative mode.")
+        if not check_hook_fn("register_forward_pre_hook", hook_fn):
             return HookHandle()
-
-        if not isinstance(hook_fn, (FunctionType, MethodType)):
-            raise TypeError(f"When using 'register_forward_pre_hook(hook_fn)', the type of 'hook_fn' must be python "
-                            f"function, but got {type(hook_fn)}.")
-        if hook_fn.__code__.co_name == "staging_specialize":
-            raise TypeError(f"Decorating hook function {hook_fn.__name__} with '@jit' is not supported.")
-
         self._enable_forward_pre_hook = True
         _pynative_executor.set_hook_changed(self)
         if not hasattr(self, '_forward_pre_hook_key'):
@@ -2036,8 +2046,8 @@ class Cell(Cell_):
             hook_fn (function): Python function. Forward hook function.
 
         Returns:
-            Handle, it is an instance of `mindspore.common.hook_handle.HookHandle` and corresponding to the `hook_fn` .
-            The handle can be used to remove the added `hook_fn` by calling `handle.remove()` .
+            A handle corresponding to the `hook_fn` . The handle can be used to remove the added `hook_fn` by calling
+            `handle.remove()` .
 
         Raises:
             TypeError: If the `hook_fn` is not a function of python.
@@ -2074,17 +2084,8 @@ class Cell(Cell_):
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32,
             value= [ 2.00000000e+00]))
         """
-        if context.get_context("mode") != context.PYNATIVE_MODE:
-            logger.warning(f"'register_forward_hook' function is only supported in pynative mode, you can use "
-                           f"context.set_context to set pynative mode.")
+        if not check_hook_fn("register_forward_hook", hook_fn):
             return HookHandle()
-
-        if not isinstance(hook_fn, (FunctionType, MethodType)):
-            raise TypeError(f"When using 'register_forward_hook(hook_fn)', the type of 'hook_fn' must be python "
-                            f"function, but got {type(hook_fn)}.")
-        if hook_fn.__code__.co_name == "staging_specialize":
-            raise TypeError(f"Decorating hook function {hook_fn.__name__} with '@jit' is not supported.")
-
         self._enable_forward_hook = True
         _pynative_executor.set_hook_changed(self)
         if not hasattr(self, '_forward_hook_key'):
@@ -2136,8 +2137,8 @@ class Cell(Cell_):
             hook_fn (function): Python function. Backward hook function.
 
         Returns:
-            Handle, it is an instance of `mindspore.common.hook_handle.HookHandle` and corresponding to the `hook_fn` .
-            The handle can be used to remove the added `hook_fn` by calling `handle.remove()` .
+            A handle corresponding to the `hook_fn` . The handle can be used to remove the added `hook_fn` by calling
+            `handle.remove()` .
 
         Raises:
             TypeError: If the `hook_fn` is not a function of python.
@@ -2172,14 +2173,8 @@ class Cell(Cell_):
             >>> print(output)
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]),)
         """
-        if context.get_context("mode") != context.PYNATIVE_MODE:
-            logger.warning(f"'register_backward_hook' function is only supported in pynative mode, you can use "
-                           f"context.set_context to set pynative mode.")
+        if not check_hook_fn("register_backward_hook", hook_fn):
             return HookHandle()
-
-        if not isinstance(hook_fn, (FunctionType, MethodType)):
-            raise TypeError(f"When using 'register_backward_hook(hook_fn)', the type of 'hook_fn' must be python "
-                            f"function, but got {type(hook_fn)}.")
         if self._cell_backward_hook is None:
             self._enable_backward_hook = True
             self._cell_backward_hook = inner.CellBackwardHook(self.cls_name + "(" + str(id(self)) + ")")
@@ -2209,10 +2204,16 @@ class Cell(Cell_):
         else:
             inputs = self._cell_backward_hook(*inputs)
             inputs = (inputs,)
-        if isinstance(inputs, tuple):
-            outputs = self.construct(*inputs, **kwargs)
+        if self.recompute_cell is not None:
+            if isinstance(inputs, tuple):
+                outputs = self.recompute_cell(*inputs, **kwargs)
+            else:
+                outputs = self.recompute_cell(inputs, **kwargs)
         else:
-            outputs = self.construct(inputs, **kwargs)
+            if isinstance(inputs, tuple):
+                outputs = self.construct(*inputs, **kwargs)
+            else:
+                outputs = self.construct(inputs, **kwargs)
         outputs = self._cell_backward_hook(outputs)
         return outputs
 
@@ -2342,6 +2343,9 @@ class Cell(Cell_):
                 introduced by optimizer shard are recomputed in auto parallel or semi auto parallel mode.
                 Default: ``False`` .
         """
+        if context.get_context("mode") == context.PYNATIVE_MODE:
+            self.recompute_cell = recompute_registry.get()(self.construct)
+            return
         self._recompute()
         if 'mp_comm_recompute' in kwargs.keys():
             self._mp_comm_recompute(kwargs.get('mp_comm_recompute', False))

mindspore/nn/extend/__init__.py

@@ -0,0 +1,29 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""
+nn Extend.
+"""
+from __future__ import absolute_import
+
+from mindspore.nn.extend.embedding import Embedding
+from mindspore.nn.extend.basic import Linear
+from mindspore.nn.extend.pooling import MaxPool2d
+from mindspore.nn.extend import layer
+from mindspore.nn.extend.layer import *
+
+__all__ = ['Embedding', 'Linear', 'MaxPool2d']
+__all__.extend(layer.__all__)
+
+__all__.sort()

mindspore/nn/extend/basic.py

@@ -0,0 +1,140 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""basic"""
+from __future__ import absolute_import
+
+import math
+
+import mindspore.common.dtype as mstype
+from mindspore import _checkparam as Validator
+from mindspore._extends import cell_attr_register
+from mindspore.common.initializer import initializer, HeUniform, Uniform
+from mindspore.common.parameter import Parameter
+from mindspore.common.tensor import Tensor
+from mindspore.nn.cell import Cell
+from mindspore.ops import operations as P
+
+__all__ = ['Linear']
+
+
+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
+
+    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`` .
+
+    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`.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor
+        >>> from mindspore.nn.extend import Linear
+        >>> import numpy as np
+        >>> x = Tensor(np.array([[180, 234, 154], [244, 48, 247]]), mindspore.float32)
+        >>> net = 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

mindspore/nn/extend/embedding.py

@@ -0,0 +1,143 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""embedding"""
+from __future__ import absolute_import
+
+import mindspore.common.dtype as mstype
+from mindspore.common.initializer import Normal
+from mindspore import _checkparam as Validator
+from mindspore.nn.cell import Cell
+from mindspore import ops
+from mindspore.common.parameter import Parameter
+from mindspore.common.tensor import Tensor
+
+__all__ = ['Embedding']
+
+
+class Embedding(Cell):
+    r"""
+    Embedding layer.
+    Retrieve the word embeddings in weight stored in the layer using indices specified in `input`.
+
+    .. warning::
+        On Ascend, the behavior is unpredictable when the value of `input` is invalid.
+
+    Args:
+        num_embeddings (int): Size of the dictionary of embeddings.
+        embedding_dim (int): The size of each embedding vector.
+        padding_idx (int, optional): If the value is not None, the corresponding row of embedding vector
+            will not be updated in training. The value of embedding vector at `padding_idx` will default
+            to zeros when the Embedding layer is newly constructed. The value should be in range
+            `[-num_embeddings, num_embeddings)` if it's not ``None``. Default ``None``.
+        max_norm (float, optional): If the value is not None, firstly get the p-norm result of the embedding
+            vector specified by `input` where p is specified by `norm_type`; if the result is larger then `max_norm`,
+            update the embedding vector` with :math:`\frac{max\_norm}{result+1e^{-7}}`. Default ``None``.
+        norm_type (float, optional): Indicated the value of p in p-norm. Default ``2.0``.
+        scale_grad_by_freq (bool, optional): If ``True`` the gradients will be scaled by the inverse of frequency
+            of the index in `input`. Default ``False``.
+        _weight (Tensor, optional): Used to initialize the weight of Embedding. If ``None``, the weight will be
+            initialized from normal distribution :math:`{N}(\text{sigma=1.0}, \text{mean=0.0})`. Default ``None``.
+        dtype (mindspore.dtype, optional) : Dtype of Parameters. It is meaningless when `_weight` is not None.
+            Default: ``mindspore.float32``.
+
+    Inputs:
+        - **input** (Tensor) - The indices used to lookup in the embedding vector. The data type must be
+          mindspore.int32 or mindspore.int64, and the value should be in range `[0, num_embeddings)`.
+
+    Outputs:
+        Tensor, has the same data type as weight, the shape is :math:`(*input.shape, embedding_dim)`.
+
+    Raises:
+        TypeError: If `num_embeddings` is not an int.
+        TypeError: If `embedding_dim` is not an int.
+        ValueError: If `padding_idx` is out of valid range.
+        TypeError: If `max_norm` is not a float.
+        TypeError: If `norm_type` is not a float.
+        TypeError: If `scale_grad_by_freq` is not a bool.
+        TypeError: If `dtype` is not one of mindspore.dtype.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> from mindspore import Tensor, nn
+        >>> input = Tensor([[1, 0, 1, 1], [0, 0, 1, 0]])
+        >>> embedding = nn.extend.Embedding(num_embeddings=10, embedding_dim=3)
+        >>> output = embedding(input)
+        >>> print(output)
+        [[[-0.0024154  -0.01203444  0.00811537]
+          [ 0.00233847 -0.00596091  0.00536799]
+          [-0.0024154  -0.01203444  0.00811537]
+          [-0.0024154  -0.01203444  0.00811537]]
+         [[ 0.00233847 -0.00596091  0.00536799]
+          [ 0.00233847 -0.00596091  0.00536799]
+          [-0.0024154  -0.01203444  0.00811537]
+          [ 0.00233847 -0.00596091  0.00536799]]]
+    """
+
+    def __init__(self, num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2.0,
+                 scale_grad_by_freq=False, _weight=None, dtype=mstype.float32):
+        """Initialize Embedding."""
+        super().__init__()
+        self.num_embeddings = Validator.check_value_type(
+            'num_embeddings', num_embeddings, [int], self.cls_name)
+        self.embedding_dim = Validator.check_value_type(
+            'embedding_dim', embedding_dim, [int], self.cls_name)
+        Validator.check_subclass(
+            "dtype", dtype, mstype.number_type, self.cls_name)
+        self.dtype = dtype
+        self.padding_idx = padding_idx
+        if _weight is None:
+            init_tensor = Tensor(shape=[num_embeddings, embedding_dim], dtype=dtype, init=Normal(1, 0))
+            init_tensor = self._zero_weight_by_index(init_tensor)
+            self.weight = Parameter(init_tensor, name='weight')
+        else:
+            self.weight = Parameter(_weight)
+
+        self.max_norm = max_norm
+        if max_norm is not None:
+            self.max_norm = Validator.check_value_type('max_norm', max_norm, [float], self.cls_name)
+
+        self.norm_type = norm_type
+        if norm_type is not None:
+            self.norm_type = Validator.check_value_type('norm_type', norm_type,
+                                                        [float], self.cls_name)
+
+        self.scale_grad_by_freq = scale_grad_by_freq
+        if scale_grad_by_freq is not None:
+            self.scale_grad_by_freq = Validator.check_value_type('scale_grad_by_freq',
+                                                                 scale_grad_by_freq,
+                                                                 [bool], self.cls_name)
+
+    def _zero_weight_by_index(self, init_tensor):
+        if self.padding_idx is not None:
+            self.padding_idx = Validator.check_int_range(self.padding_idx, -self.num_embeddings, self.num_embeddings,
+                                                         Validator.INC_LEFT, "padding_idx", self.cls_name)
+            if isinstance(init_tensor, Tensor) and init_tensor.init is not None:
+                init_tensor = init_tensor.init_data()
+            init_tensor[self.padding_idx] = 0
+
+        return init_tensor
+
+    def construct(self, input):
+        return ops.embedding(input, self.weight, self.padding_idx, self.max_norm,
+                             self.norm_type, self.scale_grad_by_freq)
+
+    def extend_repr(self):
+        return f'num_embeddings={self.num_embeddings}, embedding_dim={self.embedding_dim}, '  \
+               f'padding_idx={self.padding_idx}, max_norm={self.max_norm}, norm_type={self.norm_type}, ' \
+               f'scale_grad_by_freq={self.scale_grad_by_freq}, dtype={self.dtype}'

mindspore/nn/extend/layer/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""
+Layer.
+
+The high-level components(Cells) used to construct the neural network.
+"""
+from __future__ import absolute_import
+
+from mindspore.nn.extend.layer import normalization
+from mindspore.nn.extend.layer.normalization import *
+
+__all__ = []
+
+__all__.extend(normalization.__all__)