support create_graph in paddle.Tensor.backward() and paddle.autograd.backward()

Author: algorithm1832

paddle/fluid/eager/backward.cc

@@ -746,14 +746,15 @@ std::vector<paddle::Tensor> RunBackward(
 void Backward(const std::vector<paddle::Tensor>& tensors,  // outputs
               const std::vector<paddle::Tensor>& grad_tensors,
               bool retain_graph,
+              bool create_graph,
               std::string dump_backward_graph_path) {
   VLOG(3) << "Run in Backward";
   phi::RecordEvent backward_record_event(
       "backward", phi::TracerEventType::UserDefined, 1);
   RunBackward(tensors,
               grad_tensors,
               retain_graph,
-              false,
+              create_graph,
               {},
               false,
               {},

paddle/fluid/eager/backward.h

@@ -26,6 +26,7 @@ namespace egr {
 TEST_API void Backward(const std::vector<paddle::Tensor>& tensors,
                        const std::vector<paddle::Tensor>& grad_tensors,
                        bool retain_graph = false,
+                       bool create_graph = false,
                        std::string dump_backward_graph_path = "");
 
 TEST_API std::vector<paddle::Tensor> Grad(

paddle/fluid/pybind/eager_functions.cc

@@ -158,8 +158,9 @@ static PyObject* eager_api_run_backward(PyObject* self,
   auto tensors = CastPyArg2VectorOfTensor(PyTuple_GET_ITEM(args, 0), 0);
   auto grad_tensors = CastPyArg2VectorOfTensor(PyTuple_GET_ITEM(args, 1), 1);
   bool retain_graph = CastPyArg2AttrBoolean(PyTuple_GET_ITEM(args, 2), 2);
+  bool create_graph = CastPyArg2AttrBoolean(PyTuple_GET_ITEM(args, 3), 3);
   std::string dump_backward_graph_path =
-      CastPyArg2AttrString(PyTuple_GET_ITEM(args, 3), 3);
+      CastPyArg2AttrString(PyTuple_GET_ITEM(args, 4), 4);
   const phi::distributed::ProcessMesh* mesh = nullptr;
   if (InputsContainDistTensor(&mesh, tensors, grad_tensors)) {
     tensors = CastPyArg2VectorOfTensor(PyTuple_GET_ITEM(args, 0), 0, mesh);
@@ -168,8 +169,11 @@ static PyObject* eager_api_run_backward(PyObject* self,
   {
     eager_gil_scoped_release guard;
     EagerSetDeviceId();
-    egr::Backward(
-        tensors, grad_tensors, retain_graph, dump_backward_graph_path);
+    egr::Backward(tensors,
+                  grad_tensors,
+                  retain_graph,
+                  create_graph,
+                  dump_backward_graph_path);
   }
   RETURN_PY_NONE
   EAGER_CATCH_AND_THROW_RETURN_NULL

python/paddle/autograd/backward_mode.py

@@ -35,6 +35,7 @@ def backward(
     tensors: Tensor | Sequence[Tensor],
     grad_tensors: Tensor | Sequence[Tensor | None] | None = None,
     retain_graph: bool = False,
+    create_graph: bool = False,
     *,
     dump_backward_graph_path: str | None = None,
 ) -> None:
@@ -143,5 +144,9 @@ def check_tensors(
     assert isinstance(retain_graph, bool), "retain_graph must be True or False"
     check_and_create_dir(dump_backward_graph_path)
     core.eager.run_backward(
-        tensors, grad_tensors, retain_graph, dump_backward_graph_path
+        tensors,
+        grad_tensors,
+        retain_graph,
+        create_graph,
+        dump_backward_graph_path,
     )

python/paddle/base/dygraph/tensor_patch_methods.py

@@ -288,6 +288,7 @@ def backward(
         self: Tensor,
         grad_tensor: Tensor | None = None,
         retain_graph: bool = False,
+        create_graph: bool = False,
         *,
         dump_backward_graph_path: str | None = None,
     ) -> None:
@@ -368,7 +369,11 @@ def backward(
                 self = _grad_scalar.scale(self)
             check_and_create_dir(dump_backward_graph_path)
             core.eager.run_backward(
-                [self], grad_tensor, retain_graph, dump_backward_graph_path
+                [self],
+                grad_tensor,
+                retain_graph,
+                create_graph,
+                dump_backward_graph_path,
             )
 
             if in_profiler_mode():

test/legacy_test/test_backward.py

@@ -101,6 +101,91 @@ def test_strip_grad_suffix(self):
             self.assertEqual(backward._strip_grad_suffix_(input_), desired)
 
 
+class TestBackwardCreateGraph(unittest.TestCase):
+    """Test backward with create_graph parameter for higher-order gradients."""
+
+    def test_backward_create_graph_second_order(self):
+        """Test computing second-order gradients using create_graph=True."""
+        paddle.disable_static()
+        x = paddle.to_tensor([1.0, 2.0], dtype='float32', stop_gradient=False)
+        y = x**2
+        loss = y.sum()
+        # First backward with create_graph=True
+        # x.grad should be [2.0, 4.0]
+        paddle.autograd.backward(loss, create_graph=True)
+        grad_sum = x.grad.sum()
+        grad_sum.backward()
+        # Check second-order gradients
+        # sum up to [4.0, 6.0]
+        self.assertIsNotNone(x.grad)
+        np.testing.assert_allclose(x.grad.numpy(), [4.0, 6.0], rtol=1e-5)
+
+    def test_backward_create_graph_with_multiple_tensors(self):
+        """Test backward with create_graph on multiple output tensors."""
+        paddle.disable_static()
+        x = paddle.to_tensor(
+            [[1.0, 2.0], [3.0, 4.0]], dtype='float32', stop_gradient=False
+        )
+        z1 = x**2
+        z2 = x * 3
+        # Backward on z1
+        paddle.autograd.backward(z1, create_graph=True)
+        self.assertIsNotNone(x.grad)
+        self.assertFalse(x.grad.stop_gradient)
+        x.clear_grad()
+        # Backward on z2
+        paddle.autograd.backward(z2, create_graph=True)
+        self.assertIsNotNone(x.grad)
+        self.assertFalse(x.grad.stop_gradient)
+
+    def test_backward_create_graph_with_grad_tensors(self):
+        """Test backward with create_graph and custom grad_tensors."""
+        paddle.disable_static()
+        x = paddle.to_tensor([1.0, 2.0], dtype='float32', stop_gradient=False)
+        y = x**2
+        z = y.sum()
+        grad_tensor = paddle.to_tensor([1.0, 2.0], dtype='float32')
+        paddle.autograd.backward(y, grad_tensors=grad_tensor, create_graph=True)
+        # Check gradients with custom weights
+        self.assertIsNotNone(x.grad)
+        expected = [2.0 * 1.0, 4.0 * 2.0]  # dy/dx * weights
+        np.testing.assert_allclose(x.grad.numpy(), expected, rtol=1e-5)
+        self.assertFalse(x.grad.stop_gradient)
+
+    def test_backward_create_graph_retain_graph(self):
+        """Test backward with create_graph=True and retain_graph=True."""
+        paddle.disable_static()
+        x = paddle.to_tensor([2.0], dtype='float32', stop_gradient=False)
+        y = x**3
+        loss = y.sum()
+        # First backward
+        paddle.autograd.backward(loss, create_graph=True, retain_graph=True)
+        grad1 = x.grad.clone()
+        x.clear_grad()
+        # Second backward with same graph
+        paddle.autograd.backward(loss, create_graph=True, retain_graph=False)
+        grad2 = x.grad
+        # Gradients should be the same
+        np.testing.assert_allclose(grad1.numpy(), grad2.numpy(), rtol=1e-5)
+
+    def test_backward_create_graph_chain_rule(self):
+        """Test chain rule with higher-order gradients."""
+        paddle.disable_static()
+        x = paddle.to_tensor([1.0], dtype='float32', stop_gradient=False)
+        y = x**3
+        loss = y**2
+        # First backward
+        paddle.autograd.backward(loss, create_graph=True, retain_graph=True)
+        # At x=1: x.grad should be 6.0
+        np.testing.assert_allclose(x.grad.numpy(), [6.0], rtol=1e-5)
+        # Second backward
+        grad_sum = paddle.sum(x.grad)
+        paddle.autograd.backward(grad_sum)
+        # At x=1: x.grad should be 30
+        # Sum up to 36
+        np.testing.assert_allclose(x.grad.numpy(), [36.0], rtol=1e-5)
+
+
 if __name__ == '__main__':
     paddle.enable_static()
     unittest.main()

test/legacy_test/test_imperative_partial_backward.py

@@ -52,6 +52,33 @@ def test_partial_backward(self):
             linear1.clear_gradients()
             linear2.clear_gradients()
 
+    def test_backward_with_create_graph(self):
+        """Test backward with create_graph=True for second-order gradients"""
+        with base.dygraph.guard():
+            x = paddle.to_tensor(
+                np.array([[1.0, 2.0], [3.0, 4.0]]), dtype='float32'
+            )
+            x.stop_gradient = False
+            y = x * x
+            loss = paddle.sum(y)
+            # First backward with create_graph=True
+            loss.backward(create_graph=True)
+            # Verify first-order gradients
+            self.assertIsNotNone(x.grad)
+            first_grad = x.grad.numpy()
+            np.testing.assert_allclose(
+                first_grad, np.array([[2.0, 4.0], [6.0, 8.0]]), rtol=1e-7
+            )
+            # Compute second-order gradients
+            grad_sum = paddle.sum(x.grad)
+            grad_sum.backward()
+            # Verify second-order gradients
+            self.assertIsNotNone(x.grad)
+            second_grad = x.grad.numpy()
+            np.testing.assert_allclose(
+                second_grad, np.array([[4.0, 6.0], [8.0, 10.0]]), rtol=1e-7
+            )
+
 
 if __name__ == '__main__':
     unittest.main()