Update for scatter support fake 2d index (#47946)

paddle/phi/infermeta/ternary.cc

@@ -987,13 +987,22 @@ void ScatterInferMeta(const MetaTensor& x,
   const auto& updates_dims = updates.dims();
   const auto& ref_dims = x.dims();
   const auto& index_dims = index.dims();
-  PADDLE_ENFORCE_EQ(
-      index_dims.size(),
-      1,
-      phi::errors::InvalidArgument(
-          "The size of Input(Ids)'s shape should be equal to 1, but "
-          "received the rank of Input(Ids) is %d.",
-          index_dims.size()));
+
+  if (index_dims.size() == 2) {
+    PADDLE_ENFORCE_EQ(index_dims[1],
+                      1,
+                      phi::errors::InvalidArgument(
+                          "The last dim of the index should be 1 when the "
+                          "index is a 2D tensor, but we get %d.",
+                          index_dims[1]));
+  } else {
+    PADDLE_ENFORCE_EQ(
+        index_dims.size(),
+        1,
+        phi::errors::InvalidArgument("The index should be a 1D tensor when the "
+                                     "index is not a 2D tensor, but we get %d.",
+                                     index_dims.size()));
+  }
   PADDLE_ENFORCE_EQ(
       ref_dims.size(),
       updates_dims.size(),

python/paddle/fluid/tests/unittests/test_scatter_op.py

@@ -191,6 +191,25 @@ class TestScatterOp5(OpTest):
             )
 
 
+class TestScatterOp6(OpTest):
+    def setUp(self):
+        self.op_type = "scatter"
+        self.python_api = paddle.scatter
+        ref_np = np.ones((3, 50)).astype("float32")
+        index_np = np.array([[1], [2]]).astype("int32")
+        updates_np = np.random.random((2, 50)).astype("float32")
+        output_np = np.copy(ref_np)
+        output_np[np.array([1, 2]).astype("int32")] = updates_np
+        self.inputs = {'X': ref_np, 'Ids': index_np, 'Updates': updates_np}
+        self.outputs = {'Out': output_np}
+
+    def test_check_output(self):
+        self.check_output(check_eager=False)
+
+    def test_check_grad(self):
+        self.check_grad(["X", "Updates"], "Out", check_eager=False)
+
+
 class TestScatterAPI(unittest.TestCase):
     def setUp(self):
         self.places = [fluid.CPUPlace()]