PyTorch Backend: Missing autograd wrapper for se_t descriptor's second-order derivatives

[OutisLi]

Summary

The se_t descriptor (TabulateFusionSeTOp) does not seem to support second-order derivative calculations (required for virial/stress), even though the underlying C++/CUDA kernel (deepmd::tabulate_fusion_se_t_grad_grad_gpu) and its wrapper function (TabulateFusionSeTGradGradForward) are present in the code.

DeePMD-kit Version

devel latest

Backend and its version

pytorch

Python Version, CUDA Version, GCC Version, LAMMPS Version, etc

No response

Details

1. Working Implementation (se_a): The se_a descriptor has a complete autograd chain. Its first-derivative calculation is wrapped in TabulateFusionSeAGradOp, whose backward method correctly calls TabulateFusionSeAGradGradForward. This makes the operation fully second-order differentiable.

2. Incomplete Implementation (se_t): For the se_t descriptor, the TabulateFusionSeTOp's backward method calls TabulateFusionSeTGradForward directly. There is no corresponding TabulateFusionSeTGradOp class to wrap this first-derivative calculation.

3. The Consequence: Because the autograd wrapper (GradOp) is missing, the TabulateFusionSeTGradGradForward function is defined but is never actually called by the PyTorch autograd engine. This effectively makes it dead code in the context of PyTorch's automatic differentiation and prevents training models with virial/stress labels when using the se_t descriptor.

[iProzd]

@OutisLi The virial/stress only needs first-order derivative as same as force, while the second-order derivative calculations are for training.

@njzjz I'm also confusing that it seems that even for se_a, the TabulateFusionSeAOp in source/op/pt/tabulate_multi_device.cc only supports first-order derivative since it does not link to TabulateFusionSeAGradOp?

[njzjz]

@njzjz I'm also confusing that it seems that even for se_a, the TabulateFusionSeAOp in source/op/pt/tabulate_multi_device.cc only supports first-order derivative since it does not link to TabulateFusionSeAGradOp?

Yes, it's an incorrect implementation that doesn't support higher-order derivatives. A correct backward function refers to the code below

deepmd-kit/source/op/pt/thsilu.cc

Lines 230 to 241 in 41e62d6

	template <typename FPTYPE>
	static torch::autograd::variable_list backward_t(
	torch::autograd::AutogradContext* ctx,
	torch::autograd::variable_list grad_output) {
	// load data
	torch::autograd::variable_list saved_variables = ctx->get_saved_variables();
	torch::Tensor x_tensor = saved_variables[0];
	torch::Tensor dy_tensor = grad_output[0];
	FPTYPE w = ctx->saved_data["w"].toDouble();
	FPTYPE a = ctx->saved_data["a"].toDouble();
	return {ThsiluGradOp::apply(x_tensor, dy_tensor, w, a)[0], at::Tensor(),
	at::Tensor(), at::Tensor()};