PaddlePaddle
diff --git a/‎paddle/fluid/primitive/codegen/gen.py
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/primitive/codegen/gen.py
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/primitive/rule/vjp/details.h
Copy file name to clipboard
+8-2Lines changed: 8 additions & 2 deletions b/‎paddle/fluid/primitive/rule/vjp/details.h
Copy file name to clipboard
+8-2Lines changed: 8 additions & 2 deletions
diff --git a/‎python/paddle/autograd/backward_utils.py
Copy file name to clipboard
+6-4Lines changed: 6 additions & 4 deletions b/‎python/paddle/autograd/backward_utils.py
Copy file name to clipboard
+6-4Lines changed: 6 additions & 4 deletions
diff --git a/‎test/prim/pir_prim/CMakeLists.txt
Copy file name to clipboard
+20-6Lines changed: 20 additions & 6 deletions b/‎test/prim/pir_prim/CMakeLists.txt
Copy file name to clipboard
+20-6Lines changed: 20 additions & 6 deletions
@@ -96,6 +96,7 @@
     'atan_grad',
     'atan2_grad',
     'cumsum_grad',
+    'cumprod_grad',
     'sum_grad',
     'cast_grad',
     'reshape_grad',
 
@@ -493,8 +493,14 @@ void scatter_grad(const Tensor& index,
                   Tensor* x_grad,
                   Tensor* updates_grad) {
   if (x_grad) {
-    auto zero_tensor =
-        full<T>(common::vectorize(updates.dims()), 0.0, updates.dtype());
+    Tensor zero_tensor;
+    if (has_dynamic_shape(updates.shape())) {
+      zero_tensor =
+          backend::full_with_tensor<T>(shape<T>(updates), 0.0, updates.dtype());
+    } else {
+      zero_tensor =
+          full<T>(common::vectorize(updates.dims()), 0.0, updates.dtype());
+    }
     auto tmp_grad = scatter<T>(out_grad, index, zero_tensor, false);
     set_output<T>(tmp_grad, x_grad);
   }
 
@@ -32,7 +32,7 @@
     "pd_op.abs",
     "pd_op.add",
     "pd_op.assign",
-    "pd_op.batch_norm",
+    "pd_op.batch_norm_",
     "pd_op.cast",
     "pd_op.concat",
     "pd_op.cos",
@@ -65,20 +65,22 @@
     "pd_op.reshape",
     "pd_op.rsqrt",
     "pd_op.scale",
+    "pd_op.scatter",
+    "pd_op.scatter_nd_add",
     "pd_op.sigmoid",
     "pd_op.silu",
     "pd_op.sin",
-    "pd_op.subtract",
-    "pd_op.sum",
     "pd_op.softmax",
     "pd_op.softsign",
     "pd_op.split",
     "pd_op.sqrt",
     "pd_op.square",
     "pd_op.stack",
+    "pd_op.subtract",
+    "pd_op.sum",
     "pd_op.swiglu",
-    "pd_op.transpose",
     "pd_op.tanh",
+    "pd_op.transpose",
     "pd_op.where",
 ]
 
 
@@ -18,12 +18,6 @@ set(TEST_PRIM_PURE_PIR_CASES
     test_auto_recompute
     test_auto_recompute_dy2static
     test_prim_sub_graph_dynamic_shape
-    test_prim_sub_graph_backward_dynamic_shape
-    test_prim_sub_graph_abcde_backward_dynamic_shape
-    test_prim_sub_graph_fghij_backward_dynamic_shape
-    test_prim_sub_graph_klmno_backward_dynamic_shape
-    test_prim_sub_graph_pqrst_backward_dynamic_shape
-    test_prim_sub_graph_uvwxyz_backward_dynamic_shape
     test_decompose_control_flow
     test_decomp_whole_program
     test_dynamic_combine1
@@ -41,6 +35,26 @@ foreach(target ${TEST_PRIM_PURE_PIR_CASES})
     FLAGS_prim_enable_dynamic=true)
 endforeach()
 
+set(TEST_PRIM_DYNAMIC_SHAPE_BACKWARD_PIR_CASES
+    test_prim_sub_graph_backward_dynamic_shape
+    test_prim_sub_graph_abcde_backward_dynamic_shape
+    test_prim_sub_graph_fghij_backward_dynamic_shape
+    test_prim_sub_graph_klmno_backward_dynamic_shape
+    test_prim_sub_graph_pqrst_backward_dynamic_shape
+    test_prim_sub_graph_uvwxyz_backward_dynamic_shape)
+
+foreach(target ${TEST_PRIM_DYNAMIC_SHAPE_BACKWARD_PIR_CASES})
+  py_test_modules(
+    ${target}
+    MODULES
+    ${target}
+    ENVS
+    GLOG_v=1
+    FLAGS_prim_vjp_skip_default_ops=0
+    FLAGS_enable_pir_api=true
+    FLAGS_prim_enable_dynamic=true)
+endforeach()
+
 py_test_modules(
   test_pir_prim_flags_v3
   MODULES