A template demonstrating how to test higher-order autodiff. (#62917)

GGBond8488 · web-flow · commit 4a8ee206079c · 2024-05-10T15:00:14.000+08:00
* add composite high order test

* remove static

* add test subdirectory

* fix cmake error

* fix high grad test bug

* rm tanh

* add timeout limit

* add prim scope

* opt

* opt

* fix places

* opt
diff --git a/test/legacy_test/autodiff_checker_helper.py b/test/legacy_test/autodiff_checker_helper.py
@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
 from collections.abc import Sequence
 from logging import warning
 
 import numpy as np
+from decorator_helper import prog_scope
 
 import paddle
 from paddle import base
@@ -36,6 +38,30 @@
 }
 
 
+@contextlib.contextmanager
+def prim_guard():
+    orig_eager_prim = paddle.framework.core._is_eager_prim_enabled()
+    orig_static_prim = paddle.framework.core._is_all_prim_enabled()
+    paddle.framework.core.set_prim_eager_enabled(True)
+    paddle.framework.core._set_prim_all_enabled(True)
+    try:
+        yield
+    finally:
+        paddle.framework.core.set_prim_eager_enabled(orig_eager_prim)
+        paddle.framework.core._set_prim_all_enabled(orig_static_prim)
+
+
+def prim_scope():
+    def __impl__(fn):
+        def __fn__(*args, **kwargs):
+            with prim_guard():
+                fn(*args, **kwargs)
+
+        return __fn__
+
+    return __impl__
+
+
 def _product(t):
     return int(np.prod(t))
 
@@ -55,10 +81,12 @@ def make_jacobian(x, y_size, np_dtype):
         pass
 
 
-def compute_numerical_jacobian(program, inputs, outputs, feeds, eps):
+def compute_numerical_jacobian(program, inputs, argnums, outputs, feeds, eps):
     paddle.enable_static()
     numerical = []
-    for input in inputs:
+    for idx, input in enumerate(inputs):
+        if idx not in argnums:
+            continue
         numerical.append(
             _compute_numerical_jacobian(program, input, outputs, feeds, eps)
         )
@@ -184,19 +212,22 @@ def dtype_to_np_dtype(dtype):
         raise ValueError("Not supported data type " + str(dtype))
 
 
-def get_eager_vjp(func, inputs, cotangents=None, order=1):
+def get_eager_vjp(func, inputs, kwargs, argnums, cotangents=None, order=2):
     for x in inputs:
         x.stop_gradient = False
-    outputs = func(inputs)
-    return _get_eager_vjp(inputs, outputs, cotangents, order)
+    outputs = func(*inputs, **kwargs)
+    return _get_eager_vjp(inputs, argnums, outputs, cotangents, order)
 
 
-def _get_eager_vjp(inputs, outputs, tangents, order):
+def _get_eager_vjp(inputs, argnums, outputs, tangents, order):
     if order > 1:
         create_graph = True
     else:
         create_graph = False
 
+    if order == 1:
+        inputs = [inputs[i] for i in argnums]
+
     d_inputs = paddle.grad(
         outputs=outputs,
         inputs=inputs,
@@ -212,20 +243,21 @@ def _get_eager_vjp(inputs, outputs, tangents, order):
             ddy = paddle.ones(shape=d_input.shape, dtype=d_input.dtype)
             ddy.stop_gradient = False
             ddys.append(ddy)
-        return _get_eager_vjp(inputs, d_inputs, ddys, order - 1)
+        return _get_eager_vjp(inputs, argnums, d_inputs, ddys, order - 1)
 
     return d_inputs
 
 
 def get_static_vjp(program, feeds, fetch):
     paddle.enable_static()
     exe = paddle.static.Executor()
-    res = exe.run(program, feed=feeds, fetch_list=[fetch])
+    res = exe.run(program, feed=feeds, fetch_list=fetch)
+    res = res[: len(fetch)]
     paddle.disable_static()
     return res
 
 
-def get_static_vjp_program(func, inputs, order):
+def get_static_vjp_program(func, inputs, argnums, kwargs, order):
     cotangents = []
     paddle.enable_static()
     input_vars = []
@@ -237,7 +269,7 @@ def get_static_vjp_program(func, inputs, order):
         )
         input_vars.append(input_var)
         feeds.update({'input_' + str(idx): input.numpy()})
-    outputs = func(input_vars)
+    outputs = func(*input_vars, **kwargs)
     outputs = _as_list(outputs)
     # TODO(GGBond8488): Need to be fixed when paddle uses pir by default.
     program, (keys, values) = paddle.base.libpaddle.pir.clone_program(
@@ -264,15 +296,15 @@ def get_static_vjp_program(func, inputs, order):
             feeds.update({'grad_in_' + str(idx): grad_in_np})
             cotangents.append(grad_in_np)
         feeds, pre_outputs, d_inputs, last_grads_in = _get_static_vjp_program(
-            pir_inputs, pir_outputs, feeds, grads_in_init, order
+            pir_inputs, argnums, pir_outputs, feeds, grads_in_init, order
         )
     if not d_inputs:
-        warning(f"{func.__name__} {order}s grad will return None")
+        warning(f"{func.__name__} {order}'s grad will return None")
     paddle.disable_static()
     return program, pir_inputs, d_inputs, pre_outputs, feeds, cotangents
 
 
-def _get_static_vjp_program(inputs, outputs, feeds, grads_in, order):
+def _get_static_vjp_program(inputs, argnums, outputs, feeds, grads_in, order):
     def _require_grads(vars):
         for var in vars:
             var.stop_gradient = False
@@ -283,6 +315,8 @@ def _require_grads(vars):
     _require_grads(inputs)
     _require_grads(outputs)
     _require_grads(grads_in)
+    if order == 1:
+        inputs = [inputs[i] for i in argnums]
     d_inputs = paddle.base.gradients(outputs, inputs, grads_in)
     d_inputs = [d_input for d_input in d_inputs if d_input is not None]
     _require_grads(d_inputs)
@@ -300,12 +334,56 @@ def _require_grads(vars):
             feeds.update({f'dy_{idx}_{order}': ones})
             ddys.append(ddy)
         _require_grads(ddys)
-        return _get_static_vjp_program(inputs, d_inputs, feeds, ddys, order - 1)
+        return _get_static_vjp_program(
+            inputs, argnums, d_inputs, feeds, ddys, order - 1
+        )
     return feeds, outputs, d_inputs, grads_in
 
 
-def check_vjp(func, args, order=2, atol=None, rtol=None, eps=EPS):
+def check_vjp(
+    func,
+    args,
+    kwargs={},
+    argnums=None,
+    order=2,
+    atol=None,
+    rtol=None,
+    eps=EPS,
+):
+    """
+    Args:
+        func: function to check at ``f(args, kwargs)``.
+        args: tuple of paddle.tensor.
+        kwargs: attributes required by the func
+        order: vjp up to this order are checked.
+        atol: absolute tolerance for gradient equality.
+        rtol: relative tolerance for gradient equality.
+        eps: step size used for finite differences.
+
+    Raises:
+        AssertionError: if vjp do not match.
+    """
+    for ad_order in range(2, order + 1):
+        _check_vjp(func, args, kwargs, argnums, ad_order, atol, rtol, eps)
+
+
+@prim_scope()
+@prog_scope()
+def _check_vjp(
+    func,
+    args,
+    kwargs={},
+    argnums=None,
+    order=2,
+    atol=None,
+    rtol=None,
+    eps=EPS,
+):
     args = _as_list(args)
+    if argnums is None:
+        argnums = _as_list(range(len(args)))
+    else:
+        argnums = _as_list(argnums)
     np_type = dtype_to_np_dtype(args[0].dtype)
     atol = atol if atol else default_gradient_tolerance[np_type]
     rtol = rtol if rtol else default_gradient_tolerance[np_type]
@@ -317,18 +395,28 @@ def check_vjp(func, args, order=2, atol=None, rtol=None, eps=EPS):
         outputs,
         feeds,
         cotangents,
-    ) = get_static_vjp_program(func, args, order)
+    ) = get_static_vjp_program(func, args, argnums, kwargs, order)
     numeric_jacobian = compute_numerical_jacobian(
-        program, inputs, outputs, feeds, eps
+        program, inputs, argnums, outputs, feeds, eps
     )
     cotangents = list(map(paddle.to_tensor, cotangents))
-    eager_vjps = get_eager_vjp(func, args, cotangents, order)
+    eager_vjps = get_eager_vjp(func, args, kwargs, argnums, cotangents, order)
     static_vjps_np = get_static_vjp(program, feeds, fetch_list)
     eager_vjps_np = []
     for eager_vjp in eager_vjps:
         eager_vjps_np.append(eager_vjp.numpy())
     inputs_length = len(numeric_jacobian)
     numeric_vjps = []
+    if not eager_vjps_np:
+        if not static_vjps_np:
+            warning(
+                f"{func.__name__} in {order}'s derivative will return zeros(None)"
+            )
+            return
+        raise ValueError("eager returns None vjp, but static not")
+    if not static_vjps_np:
+        static_vjps_np = np.zeros_like(numeric_vjps)
+
     for x_idx in range(inputs_length):
         jacobians = _as_list(numeric_jacobian[x_idx])
         dx_idx = None
@@ -341,7 +429,6 @@ def check_vjp(func, args, order=2, atol=None, rtol=None, eps=EPS):
         numeric_vjps.append(dx_idx)
     eager_vjps_np = list(map(np.ndarray.flatten, eager_vjps_np))
     static_vjps_np = list(map(np.ndarray.flatten, static_vjps_np))
-
     np.testing.assert_allclose(
         numeric_vjps,
         eager_vjps_np,
diff --git a/test/prim/CMakeLists.txt b/test/prim/CMakeLists.txt
@@ -13,3 +13,4 @@ add_subdirectory(model)
 add_subdirectory(composite_ops)
 add_subdirectory(process)
 add_subdirectory(pir_prim)
+add_subdirectory(high_order)
diff --git a/test/prim/high_order/CMakeLists.txt b/test/prim/high_order/CMakeLists.txt
@@ -0,0 +1,12 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS FLAGS_enable_pir_api=true
+                  FLAGS_prim_enable_dynamic=true)
+endforeach()
+
+set_tests_properties(test_high_order_derivative PROPERTIES TIMEOUT 90)
diff --git a/test/prim/high_order/test_high_order_derivative.py b/test/prim/high_order/test_high_order_derivative.py
@@ -0,0 +1,123 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+import unittest
+
+import numpy as np
+
+sys.path.append("../../legacy_test")
+
+import autodiff_checker_helper as ad_checker
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+class TestMulHigherOrderAD(unittest.TestCase):
+    @param.parameterized.expand(
+        [
+            (
+                paddle.multiply,
+                [2, 3, 2],
+                [2, 3, 2],
+                'float32',
+                2,
+                ('cpu', 'gpu'),
+            ),
+            (
+                paddle.multiply,
+                [2, 3, 2],
+                [2, 3, 2],
+                'float64',
+                2,
+                ('cpu', 'gpu'),
+            ),
+            (
+                paddle.multiply,
+                [2, 3, 2, 1],
+                [2, 3, 2, 4],
+                'float64',
+                2,
+                ('cpu', 'gpu'),
+            ),
+        ]
+    )
+    def test_high_order_autodiff(
+        self, func, shape1, shape2, dtype, order, places
+    ):
+        var_1 = np.random.randn(*shape1).astype(dtype)
+        var_2 = np.random.randn(*shape2).astype(dtype)
+        for place in places:
+            if place == 'gpu' and not core.is_compiled_with_cuda():
+                continue
+            var1 = paddle.to_tensor(var_1, place=place)
+            var2 = paddle.to_tensor(var_2, place=place)
+            var1, var2 = paddle.broadcast_tensors(input=[var1, var2])
+            ad_checker.check_vjp(
+                func, [var1, var2], argnums=(0, 1), order=order
+            )
+
+
+class TestSinHigherOrderAD(unittest.TestCase):
+    @param.parameterized.expand(
+        [
+            (paddle.sin, [2, 3, 2], 'float32', 4, ('cpu', 'gpu')),
+            (paddle.sin, [2, 3, 2, 4], 'float64', 4, ('cpu', 'gpu')),
+        ]
+    )
+    def test_high_order_autodiff(self, func, shape, dtype, order, places):
+        var_1 = np.random.randn(*shape).astype(dtype)
+        for place in places:
+            if place == 'gpu' and not core.is_compiled_with_cuda():
+                continue
+            var1 = paddle.to_tensor(var_1, place=place)
+            ad_checker.check_vjp(func, [var1], argnums=(0), order=order)
+
+
+class TestCosHigherOrderAD(TestSinHigherOrderAD):
+    @param.parameterized.expand(
+        [
+            (paddle.cos, [2, 3, 2], 'float32', 4, ('cpu', 'gpu')),
+            (paddle.cos, [2, 3, 2, 4], 'float64', 4, ('cpu', 'gpu')),
+        ]
+    )
+    def test_high_order_autodiff(self, func, shape, dtype, order, places):
+        var_1 = np.random.randn(*shape).astype(dtype)
+        for place in places:
+            if place == 'gpu' and not core.is_compiled_with_cuda():
+                continue
+            var1 = paddle.to_tensor(var_1, place=place)
+            ad_checker.check_vjp(func, [var1], argnums=(0), order=order)
+
+
+class TestTanhHigherOrderAD(TestSinHigherOrderAD):
+    @param.parameterized.expand(
+        [
+            (paddle.tanh, [2, 3, 2], 'float32', 4, ('cpu', 'gpu')),
+            (paddle.tanh, [2, 3, 2, 4], 'float64', 4, ('cpu', 'gpu')),
+        ]
+    )
+    def test_high_order_autodiff(self, func, shape, dtype, order, places):
+        var_1 = np.random.randn(*shape).astype(dtype)
+        for place in places:
+            if place == 'gpu' and not core.is_compiled_with_cuda():
+                continue
+            var1 = paddle.to_tensor(var_1, place=place)
+            ad_checker.check_vjp(func, [var1], argnums=(0), order=order)
+
+
+if __name__ == "__main__":
+    unittest.main()
