fix(pt): enable second-order autograd for tabulate descriptors

source/tests/pt/tabulate_test_utils.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from collections.abc import (
+    Sequence,
+)
+
+import torch
+
+
+def _nonuniform_like(tensor: torch.Tensor) -> torch.Tensor:
+    """Create deterministic nonuniform weights for gradient projections."""
+    return torch.linspace(
+        -0.5,
+        0.75,
+        tensor.numel(),
+        dtype=tensor.dtype,
+        device=tensor.device,
+    ).reshape_as(tensor)
+
+
+def _project_first_grads(
+    descriptor_tensor: torch.Tensor,
+    grad_inputs: Sequence[torch.Tensor],
+    dy_tensor: torch.Tensor,
+    weights: Sequence[torch.Tensor],
+) -> torch.Tensor:
+    first_grads = torch.autograd.grad(
+        descriptor_tensor,
+        grad_inputs,
+        grad_outputs=dy_tensor,
+        retain_graph=True,
+    )
+    projection = descriptor_tensor.new_zeros(())
+    for grad, weight in zip(first_grads, weights, strict=True):
+        projection = projection + (grad * weight).sum()
+    return projection
+
+
+def assert_second_order_backward_matches_finite_difference(
+    descriptor_tensor: torch.Tensor,
+    grad_inputs: torch.Tensor | Sequence[torch.Tensor],
+) -> None:
+    """Compare dz/dy from the grad-grad kernel with a finite-difference reference."""
+    if isinstance(grad_inputs, torch.Tensor):
+        grad_inputs = (grad_inputs,)
+
+    dy_tensor = _nonuniform_like(descriptor_tensor).requires_grad_(True)
+    first_grads = torch.autograd.grad(
+        descriptor_tensor,
+        grad_inputs,
+        grad_outputs=dy_tensor,
+        create_graph=True,
+        retain_graph=True,
+    )
+    weights = tuple(_nonuniform_like(grad) for grad in first_grads)
+    projection = descriptor_tensor.new_zeros(())
+    for grad, weight in zip(first_grads, weights, strict=True):
+        projection = projection + (grad * weight).sum()
+
+    dz_dy_tensor = torch.autograd.grad(
+        projection,
+        dy_tensor,
+        retain_graph=True,
+    )[0]
+
+    eps, atol, rtol = (
+        (1e-6, 1e-6, 1e-6)
+        if descriptor_tensor.dtype == torch.float64
+        else (1e-2, 5e-3, 5e-3)
+    )
+    reference = torch.empty_like(dy_tensor)
+    base_dy = dy_tensor.detach()
+    reference_flat = reference.view(-1)
+    for idx in range(dy_tensor.numel()):
+        dy_plus = base_dy.clone()
+        dy_minus = base_dy.clone()
+        dy_plus.view(-1)[idx] += eps
+        dy_minus.view(-1)[idx] -= eps
+
+        plus_projection = _project_first_grads(
+            descriptor_tensor, grad_inputs, dy_plus, weights
+        )
+        minus_projection = _project_first_grads(
+            descriptor_tensor, grad_inputs, dy_minus, weights
+        )
+        reference_flat[idx] = (plus_projection - minus_projection) / (2.0 * eps)
+
+    torch.testing.assert_close(
+        dz_dy_tensor,
+        reference,
+        atol=atol,
+        rtol=rtol,
+    )

source/tests/pt/test_tabulate_fusion_se_a.py

@@ -13,6 +13,9 @@
 from ..consistent.common import (
     parameterized,
 )
+from .tabulate_test_utils import (
+    assert_second_order_backward_matches_finite_difference,
+)
 
 
 @parameterized((torch.float64, torch.float32))
@@ -1488,6 +1491,21 @@ def test_backward(self) -> None:
             rtol=self.prec,
         )
 
+    def test_second_order_backward(self) -> None:
+        forward_result = torch.ops.deepmd.tabulate_fusion_se_a(
+            self.table_tensor,
+            self.table_info_tensor,
+            self.em_x_tensor,
+            self.em_tensor,
+            self.last_layer_size,
+        )
+
+        descriptor_tensor = forward_result[0]
+        assert_second_order_backward_matches_finite_difference(
+            descriptor_tensor,
+            (self.em_x_tensor, self.em_tensor),
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

source/tests/pt/test_tabulate_fusion_se_atten.py

@@ -13,6 +13,9 @@
 from ..consistent.common import (
     parameterized,
 )
+from .tabulate_test_utils import (
+    assert_second_order_backward_matches_finite_difference,
+)
 
 
 @parameterized((torch.float64, torch.float32))
@@ -1624,6 +1627,23 @@ def test_backward(self) -> None:
             self.em_tensor.grad, self.expected_dy_dem, atol=self.prec, rtol=self.prec
         )
 
+    def test_second_order_backward(self) -> None:
+        forward_result = torch.ops.deepmd.tabulate_fusion_se_atten(
+            self.table_tensor,
+            self.table_info_tensor,
+            self.em_x_tensor,
+            self.em_tensor,
+            self.two_embed_tensor,
+            self.last_layer_size,
+            self.is_sorted,
+        )
+
+        descriptor_tensor = forward_result[0]
+        assert_second_order_backward_matches_finite_difference(
+            descriptor_tensor,
+            (self.em_x_tensor, self.em_tensor, self.two_embed_tensor),
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

source/tests/pt/test_tabulate_fusion_se_r.py

@@ -13,6 +13,9 @@
 from ..consistent.common import (
     parameterized,
 )
+from .tabulate_test_utils import (
+    assert_second_order_backward_matches_finite_difference,
+)
 
 
 @parameterized((torch.float64, torch.float32))
@@ -1332,6 +1335,20 @@ def test_backward(self) -> None:
             self.em_tensor.grad, self.expected_dy_dem, atol=self.prec, rtol=self.prec
         )
 
+    def test_second_order_backward(self) -> None:
+        forward_result = torch.ops.deepmd.tabulate_fusion_se_r(
+            self.table_tensor,
+            self.table_info_tensor,
+            self.em_tensor,
+            self.last_layer_size,
+        )
+
+        descriptor_tensor = forward_result[0]
+        assert_second_order_backward_matches_finite_difference(
+            descriptor_tensor,
+            self.em_tensor,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

source/tests/pt/test_tabulate_fusion_se_t.py

@@ -13,6 +13,9 @@
 from ..consistent.common import (
     parameterized,
 )
+from .tabulate_test_utils import (
+    assert_second_order_backward_matches_finite_difference,
+)
 
 
 @parameterized((torch.float64, torch.float32))
@@ -1748,6 +1751,21 @@ def test_backward(self) -> None:
             rtol=self.prec,
         )
 
+    def test_second_order_backward(self) -> None:
+        forward_result = torch.ops.deepmd.tabulate_fusion_se_t(
+            self.table_tensor,
+            self.table_info_tensor,
+            self.em_x_tensor,
+            self.em_tensor,
+            self.last_layer_size,
+        )
+
+        descriptor_tensor = forward_result[0]
+        assert_second_order_backward_matches_finite_difference(
+            descriptor_tensor,
+            (self.em_x_tensor, self.em_tensor),
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

source/tests/pt/test_tabulate_fusion_se_t_tebd.py

@@ -13,6 +13,9 @@
 from ..consistent.common import (
     parameterized,
 )
+from .tabulate_test_utils import (
+    assert_second_order_backward_matches_finite_difference,
+)
 
 
 @parameterized((torch.float64, torch.float32))
@@ -1913,6 +1916,21 @@ def test_backward(self) -> None:
             rtol=self.prec,
         )
 
+    def test_second_order_backward(self) -> None:
+        forward_result = torch.ops.deepmd.tabulate_fusion_se_t_tebd(
+            self.table_tensor,
+            self.table_info_tensor,
+            self.em_x_tensor,
+            self.em_tensor,
+            self.last_layer_size,
+        )
+
+        descriptor_tensor = forward_result[0]
+        assert_second_order_backward_matches_finite_difference(
+            descriptor_tensor,
+            self.em_x_tensor,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()