deepmodeling / deepmodeling/deepmd-kit

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

Open
#5,666 0 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

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:

https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/deepmd/infer/deep_eval.py#L129-L139

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:

https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/deepmd/pt/infer/deep_eval.py#L519-L529

But the normal and spin paths reshape directly:

https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/deepmd/pt/infer/deep_eval.py#L693-L701

https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/deepmd/pt/infer/deep_eval.py#L859-L867

JAX has the same public contract and direct reshape:

https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/deepmd/jax/infer/deep_eval.py#L216-L226

https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/deepmd/jax/infer/deep_eval.py#L356-L369

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

https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/deepmd/tf2/infer/deep_eval.py#L324-L337

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

https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/deepmd/pd/infer/deep_eval.py#L351-L361

https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/deepmd/pd/infer/deep_eval.py#L542-L550

https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/deepmd/pd/infer/deep_eval.py#L635-L643

TF1 already expands these shorthand forms before evaluation:

https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/deepmd/tf/infer/deep_eval.py#L811-L829

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.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.