[Code scan] Preserve Paddle fparam/aparam inputs when freezing models

[njzjz]

This issue comes from a Codex global scan of deepmodeling/deepmd-kit at commit 73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.

Problem

The Paddle freeze entrypoint converts model.forward to static with fparam and aparam fixed to None:

deepmd-kit/deepmd/pd/entrypoints/main.py

Lines 386 to 397 in 73de44b

	if hasattr(model, "forward"):
	model.forward = paddle.jit.to_static(
	model.forward,
	input_spec=[
	InputSpec([-1, -1, 3], dtype="float64", name="coord"), # coord
	InputSpec([-1, -1], dtype="int64", name="atype"), # atype
	InputSpec([-1, 9], dtype="float64", name="box"), # box
	None, # fparam
	None, # aparam
	# InputSpec([], dtype="bool", name="do_atomic_virial"), # do_atomic_virial
	do_atomic_virial, # do_atomic_virial
	],

It does the same for model.forward_lower:

deepmd-kit/deepmd/pd/entrypoints/main.py

Lines 407 to 418 in 73de44b

	if hasattr(model, "forward_lower"):
	model.forward_lower = paddle.jit.to_static(
	model.forward_lower,
	input_spec=[
	InputSpec([-1, -1, 3], dtype="float64", name="coord"), # extended_coord
	InputSpec([-1, -1], dtype="int32", name="atype"), # extended_atype
	InputSpec([-1, -1, -1], dtype="int32", name="nlist"), # nlist
	InputSpec([-1, -1], dtype="int64", name="mapping"), # mapping
	None, # fparam
	None, # aparam
	# InputSpec([], dtype="bool", name="do_atomic_virial"), # do_atomic_virial
	do_atomic_virial, # do_atomic_virial

The lower-level serializer repeats the same static signature:

deepmd-kit/deepmd/pd/utils/serialization.py

Lines 68 to 79 in 73de44b

	model.forward = paddle.jit.to_static(
	model.forward,
	full_graph=True,
	input_spec=[
	InputSpec([-1, -1, 3], dtype="float64", name="coord"),
	InputSpec([-1, -1], dtype="int64", name="atype"),
	InputSpec([-1, 9], dtype="float64", name="box"),
	None,
	None,
	True,
	],
	)

deepmd-kit/deepmd/pd/utils/serialization.py

Lines 87 to 99 in 73de44b

	model.forward_lower = paddle.jit.to_static(
	model.forward_lower,
	full_graph=True,
	input_spec=[
	InputSpec([-1, -1, 3], dtype="float64", name="coord"),
	InputSpec([-1, -1], dtype="int32", name="atype"),
	InputSpec([-1, -1, -1], dtype="int32", name="nlist"),
	None,
	None,
	None,
	True,
	None,
	],

The model interfaces otherwise support forwarding frame and atomic parameters:

deepmd-kit/deepmd/pd/model/model/make_model.py

Lines 131 to 139 in 73de44b

	def forward_common(
	self,
	coord: paddle.Tensor,
	atype: paddle.Tensor,
	box: paddle.Tensor | None = None,
	fparam: paddle.Tensor | None = None,
	aparam: paddle.Tensor | None = None,
	do_atomic_virial: bool = False,
	) -> dict[str, paddle.Tensor]:

deepmd-kit/deepmd/pd/model/model/make_model.py

Lines 253 to 264 in 73de44b

	def forward_common_lower(
	self,
	extended_coord: paddle.Tensor,
	extended_atype: paddle.Tensor,
	nlist: paddle.Tensor,
	mapping: paddle.Tensor | None = None,
	fparam: paddle.Tensor | None = None,
	aparam: paddle.Tensor | None = None,
	do_atomic_virial: bool = False,
	comm_dict: list[paddle.Tensor] | None = None,
	extra_nlist_sort: bool = False,
	) -> dict[str, paddle.Tensor]:

and fitting nets assert that fparam/aparam are present when configured:

deepmd-kit/deepmd/pd/model/task/fitting.py

Lines 595 to 609 in 73de44b

	if self.numb_fparam > 0:
	assert fparam is not None, "fparam should not be None"
	assert self.fparam_avg is not None
	assert self.fparam_inv_std is not None
	if fparam.shape[-1] != self.numb_fparam:
	raise ValueError(
	"get an input fparam of dim {fparam.shape[-1]}, ",
	"which is not consistent with {self.numb_fparam}.",
	)
	fparam = fparam.reshape([nf, self.numb_fparam])
	nb, _ = fparam.shape
	t_fparam_avg = self._extend_f_avg_std(self.fparam_avg, nb)
	t_fparam_inv_std = self._extend_f_avg_std(self.fparam_inv_std, nb)
	fparam = (fparam - t_fparam_avg) * t_fparam_inv_std
	fparam = paddle.tile(fparam.reshape([nf, 1, -1]), [1, nloc, 1])

deepmd-kit/deepmd/pd/model/task/fitting.py

Lines 619 to 620 in 73de44b

	# check aparam dim, concate to input descriptor
	if self.numb_aparam > 0 and not self.use_aparam_as_mask:

Impact

Paddle models trained with required frame parameters or atomic parameters cannot be exported to a frozen static model that accepts those inputs. The static graph signature bakes both values as None, so inference through the frozen Paddle export cannot supply required fparam or aparam.

Suggested fix

When the loaded model reports nonzero get_dim_fparam() or get_dim_aparam(), include corresponding InputSpecs in the Paddle static signatures for both forward and forward_lower. Keep None only for models that do not use those inputs.