PaddlePaddle
diff --git a/‎paddle/fluid/primitive/codegen/gen.py
Copy file name to clipboard
+16-2Lines changed: 16 additions & 2 deletions b/‎paddle/fluid/primitive/codegen/gen.py
Copy file name to clipboard
+16-2Lines changed: 16 additions & 2 deletions
diff --git a/‎paddle/fluid/primitive/codegen/templates/rule/vjp/generated/generated_vjp.cc.j2
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/fluid/primitive/codegen/templates/rule/vjp/generated/generated_vjp.cc.j2
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/fluid/primitive/primitive.yaml
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/primitive/primitive.yaml
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
@@ -72,8 +72,20 @@
 ]
 
 
-PRIM_VJP = ['divide_grad', 'sum_grad']  # vjp list of primitive op
-CUSTOM_VJP = ['gelu_grad']  # custom vjp list of composite op
+PRIM_VJP = [
+    'divide_grad',
+    'sum_grad',
+    'cast_grad',
+    'add_grad',
+    'multiply_grad',
+    'elementwise_pow_grad',
+    'reshape_grad',
+    'split_grad',
+    'tanh_grad',
+    'transpose_grad',
+    'concat_grad',
+]  # vjp list of primitive op
+CUSTOM_VJP = ['gelu_grad', 'layer_norm_grad']  # custom vjp list of composite op
 VJP_COMPS = PRIM_VJP + CUSTOM_VJP
 
 BACKENDS = [
@@ -149,6 +161,8 @@
     'embedding_grad',
     'sqrt',
     'uniform',
+    'split',
+    'transpose',
 ]
 
 
 
@@ -106,7 +106,7 @@ paddle::Tensor* {{api.outputs[i].name}} = !stop_gradients[{{i}}][0] ? &vjp_res[{
     {% else %}
 std::vector<paddle::Tensor*> {{api.outputs[i].name}}(stop_gradients[{{i}}].size(), nullptr);
 for (size_t i=0; i< stop_gradients[{{i}}].size(); i++ ) {
-  {{api.outputs[i].name}} =  !stop_gradients[{{i}}][i] ?  &vjp_res[{{i}}][i] : nullptr;
+  {{api.outputs[i].name}}[i] =  !stop_gradients[{{i}}][i] ?  &vjp_res[{{i}}][i] : nullptr;
 }
     {% endif %}
   {% endfor %}
 
@@ -49,3 +49,4 @@
 - erf
 - tanh
 - full
+- cast
Original file line number	Diff line number	Diff line change
`@@ -106,7 +106,7 @@ paddle::Tensor* {{api.outputs[i].name}} = !stop_gradients[{{i}}][0] ? &vjp_res[{`
`106`	`106`	`{% else %}`
`107`	`107`	`std::vector<paddle::Tensor*> {{api.outputs[i].name}}(stop_gradients[{{i}}].size(), nullptr);`
`108`	`108`	`for (size_t i=0; i< stop_gradients[{{i}}].size(); i++ ) {`
`109`		`- {{api.outputs[i].name}} = !stop_gradients[{{i}}][i] ? &vjp_res[{{i}}][i] : nullptr;`
	`109`	`+ {{api.outputs[i].name}}[i] = !stop_gradients[{{i}}][i] ? &vjp_res[{{i}}][i] : nullptr;`
`110`	`110`	`}`
`111`	`111`	`{% endif %}`
`112`	`112`	`{% endfor %}`
-Original file line number
+Diff line change
 - erf
 - tanh
 - full
 +- cast