support embedding for 3d

mindspore/nn/layer/embedding.py

@@ -59,7 +59,7 @@ class Embedding(Cell):
 
     Args:
         vocab_size (int): Size of the dictionary of embeddings.
-        embedding_size (int): The size of each embedding vector.
+        embedding_size (Union[int, tuple(int), list(int)]): The size of each embedding vector.
         use_one_hot (bool): Specifies whether to apply one_hot encoding form. Default: False.
         embedding_table (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the embedding_table.
             Refer to class `initializer` for the values of string when a string
@@ -98,12 +98,21 @@ class Embedding(Cell):
                  dtype=mstype.float32, padding_idx=None):
         super(Embedding, self).__init__()
         self.vocab_size = validator.check_value_type('vocab_size', vocab_size, [int], self.cls_name)
-        self.embedding_size = validator.check_value_type('embedding_size', embedding_size, [int], self.cls_name)
+        self.embedding_size = validator.check_value_type('embedding_size', embedding_size,
+                                                         [int, tuple, list], self.cls_name)
         validator.check_value_type('use_one_hot', use_one_hot, [bool], self.cls_name)
         validator.check_subclass("dtype", dtype, mstype.number_type, self.cls_name)
         self.use_one_hot = use_one_hot
         self.dtype = dtype
-        self.init_tensor = initializer(embedding_table, [vocab_size, embedding_size])
+        if isinstance(self.embedding_size, int):
+            self.init_tensor = initializer(embedding_table, [vocab_size, embedding_size])
+            self.embedding_out = (self.embedding_size,)
+        else:
+            if len(self.embedding_size) != 2:
+                raise ValueError("embedding_size should be a int or a tuple of two ints")
+            self.init_tensor = initializer(embedding_table, [vocab_size, self.embedding_size[0],
+                                                             self.embedding_size[1]])
+            self.embedding_out = (self.embedding_size[0], self.embedding_size[1],)
         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,
@@ -127,7 +136,7 @@ class Embedding(Cell):
 
     def construct(self, ids):
         extended_ids = self.expand(ids, -1)
-        out_shape = self.get_shp(ids) + (self.embedding_size,)
+        out_shape = self.get_shp(ids) + self.embedding_out
         flat_ids = self.reshape_flat(extended_ids, self.shp_flat)
 
         if self.use_one_hot:

tests/ut/python/ops/test_ops.py

@@ -891,6 +891,7 @@ class StridedSliceNet(nn.Cell):
         out_3 = self.strided_slice_3(x, self.begins, self.ends, self.strides) + self.const_3
         return out_0, out_1, out_2, out_3
 
+
 @pytest.mark.skip(reason='0 in shape is not support')
 def test_strided_slice_const():
     class StridedSLiceConstNet(nn.Cell):
@@ -1290,6 +1291,10 @@ test_case_math_ops = [
         'block': Embedding(vocab_size=10, embedding_size=3, padding_idx=2),
         'desc_inputs': [Tensor(np.array([0, 2, 2, 7]).astype(np.int32))],
         'skip': ['backward']}),
+    ('Embedding_3', {
+        'block': Embedding(vocab_size=10, embedding_size=(4, 4), padding_idx=2),
+        'desc_inputs': [Tensor(np.random.randint(6, size=(4, 4)))],
+        'skip': ['backward']}),
     ('EmbeddingLookup_1', {
         'block': EmbeddingLookup(vocab_size=10, embedding_size=3),
         'desc_inputs': [Tensor(np.array([0, 2, 2, 7]).astype(np.int32))],
@@ -2289,7 +2294,7 @@ test_case_array_ops = [
         'desc_inputs': [Tensor(np.array([1], np.float32)),
                         Tensor(np.array([1], np.float32)),
                         Tensor(np.array([1], np.float32))],
-        'desc_bprop': [[3,]]}),
+        'desc_bprop': [[3, ]]}),
     ('Stack_0', {
         'block': NetForStackInput(P.Stack()),
         'desc_inputs': [[2, 2], [2, 2], [2, 2]],
@@ -2711,7 +2716,7 @@ test_case_other_ops = [
                         Tensor(np.random.rand(1, 64).astype(np.float16)),
                         Tensor(np.random.rand(1, 64).astype(np.float16)),
                         Tensor(np.random.rand(96, 256).astype(np.float16)),
-                        Tensor(np.random.rand(256,).astype(np.float16))],
+                        Tensor(np.random.rand(256, ).astype(np.float16))],
         'desc_bprop': [Tensor(np.random.rand(1, 64).astype(np.float16)),
                        Tensor(np.random.rand(1, 64).astype(np.float16)),
                        Tensor(np.random.rand(1, 64).astype(np.float16)),