Skip to content

[Code scan] Honor documented fparam/aparam shorthand in non-TF1 DeepEval backends #5666

Description

@njzjz

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

Problem

The public DeepEval interface documents shorthand inputs for frame and atomic parameters:

fparam
The frame parameter.
The array can be of size :
- nframes x dim_fparam.
- dim_fparam. Then all frames are assumed to be provided with the same fparam.
aparam
The atomic parameter
The array can be of size :
- nframes x natoms x dim_aparam.
- natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam.
- dim_aparam. Then all frames and atoms are provided with the same aparam.

Several non-TF1 backends implement their own inference adapters but reshape directly to the fully batched shapes. This rejects the documented shorthand whenever nframes > 1, and rejects scalar-per-atom aparam shorthand whenever natoms > 1.

PyTorch documents the shorthand here:

fparam
The frame parameter.
The array can be of size :
- nframes x dim_fparam.
- dim_fparam. Then all frames are assumed to be provided with the same fparam.
aparam
The atomic parameter
The array can be of size :
- nframes x natoms x dim_aparam.
- natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam.
- dim_aparam. Then all frames and atoms are provided with the same aparam.

But the normal and spin paths reshape directly:

if fparam is not None:
fparam_input = to_torch_tensor(
fparam.reshape(nframes, self.get_dim_fparam())
)
else:
fparam_input = None
if aparam is not None:
aparam_input = to_torch_tensor(
aparam.reshape(nframes, natoms, self.get_dim_aparam())

if fparam is not None:
fparam_input = to_torch_tensor(
fparam.reshape(nframes, self.get_dim_fparam())
)
else:
fparam_input = None
if aparam is not None:
aparam_input = to_torch_tensor(
aparam.reshape(nframes, natoms, self.get_dim_aparam())

JAX has the same public contract and direct reshape:

fparam
The frame parameter.
The array can be of size :
- nframes x dim_fparam.
- dim_fparam. Then all frames are assumed to be provided with the same fparam.
aparam
The atomic parameter
The array can be of size :
- nframes x natoms x dim_aparam.
- natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam.
- dim_aparam. Then all frames and atoms are provided with the same aparam.

if fparam is not None:
fparam_input = fparam.reshape(nframes, self.get_dim_fparam())
elif self.dp.has_default_fparam():
# JAX (XLA) requires static shapes, so default must be implemented here
default_fparam = self.dp.get_default_fparam()
assert default_fparam is not None
fparam_input = np.tile(
np.array(default_fparam, dtype=GLOBAL_NP_FLOAT_PRECISION),
(nframes, 1),
)
else:
fparam_input = None
if aparam is not None:
aparam_input = aparam.reshape(nframes, natoms, self.get_dim_aparam())

TensorFlow 2 implements the same backend interface but also reshapes directly:

if fparam is not None:
fparam_input = fparam.reshape(nframes, self.get_dim_fparam())
elif self.dp.has_default_fparam():
default_fparam = self.dp.get_default_fparam()
assert default_fparam is not None
fparam_input = np.tile(
np.array(default_fparam, dtype=GLOBAL_NP_FLOAT_PRECISION),
(nframes, 1),
)
else:
fparam_input = None
aparam_input = (
aparam.reshape(nframes, natoms, self.get_dim_aparam())
if aparam is not None

Paddle documents the shorthand but reshapes directly in both normal and spin paths:

fparam
The frame parameter.
The array can be of size :
- nframes x dim_fparam.
- dim_fparam. Then all frames are assumed to be provided with the same fparam.
aparam
The atomic parameter
The array can be of size :
- nframes x natoms x dim_aparam.
- natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam.
- dim_aparam. Then all frames and atoms are provided with the same aparam.

if fparam is not None:
fparam_input = to_paddle_tensor(
fparam.reshape([nframes, self.get_dim_fparam()])
)
else:
fparam_input = None
if aparam is not None:
aparam_input = to_paddle_tensor(
aparam.reshape([nframes, natoms, self.get_dim_aparam()])

if fparam is not None:
fparam_input = to_paddle_tensor(
fparam.reshape([nframes, self.get_dim_fparam()])
)
else:
fparam_input = None
if aparam is not None:
aparam_input = to_paddle_tensor(
aparam.reshape([nframes, natoms, self.get_dim_aparam()])

TF1 already expands these shorthand forms before evaluation:

if self.has_fparam:
fdim = self.get_dim_fparam()
if fparam.size == nframes * fdim:
fparam = np.reshape(fparam, [nframes, fdim])
elif fparam.size == fdim:
fparam = np.tile(fparam.reshape([-1]), [nframes, 1])
else:
raise RuntimeError(
f"got wrong size of frame param, should be either {nframes} x {fdim} or {fdim}"
)
if self.has_aparam:
fdim = self.get_dim_aparam()
if aparam.size == nframes * natoms * fdim:
aparam = np.reshape(aparam, [nframes, natoms * fdim])
elif aparam.size == natoms * fdim:
aparam = np.tile(aparam.reshape([-1]), [nframes, 1])
elif aparam.size == fdim:
aparam = np.tile(aparam.reshape([-1]), [nframes, natoms])
else:

This is related to, but separate from, #5662, which covers the dpmodel backend implementation.

Impact

Users who pass documented fparam.shape == (dim_fparam,), aparam.shape == (natoms, dim_aparam), or aparam.shape == (dim_aparam,) get reshape failures on multi-frame or multi-atom inference in PyTorch, JAX, TF2, and Paddle backends.

Suggested fix

Move the TF1-style parameter normalization into a shared helper and call it from every backend-level DeepEval implementation before converting to backend tensors. Add parity tests that evaluate the same model with full-shape and shorthand fparam/aparam inputs across these backends.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    Status
    In Progress

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions