Skip to content

[Code scan] Do not silently ignore TF2 DeepEval neighbor_list #5667

Description

@njzjz

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

Problem

deepmd.tf2.infer.DeepEval accepts a public neighbor_list argument:

def __init__(
self,
model_file: str,
output_def: ModelOutputDef,
*args: Any,
auto_batch_size: bool | int | AutoBatchSize = True,
neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
**kwargs: Any,

But the initializer never stores or validates it, and evaluation always calls the SavedModel wrapper without forwarding any neighbor-list object:

self.output_def = output_def
self.model_path = model_file
self.dp = TF2SavedModelWrapper(model_file)
self.rcut = self.dp.get_rcut()
self.type_map = self.dp.get_type_map()
if isinstance(auto_batch_size, bool):
self.auto_batch_size = AutoBatchSize() if auto_batch_size else None
elif isinstance(auto_batch_size, int):
self.auto_batch_size = AutoBatchSize(auto_batch_size)
elif isinstance(auto_batch_size, AutoBatchSize):
self.auto_batch_size = auto_batch_size
else:
raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize")

batch_output = self.dp(
coord_input,
type_input,
box=box_input,
fparam=fparam_input,
aparam=aparam_input,
do_atomic_virial=do_atomic_virial,

By contrast, the TF1 inference adapter stores neighbor_list and switches to build_neighbor_list when it is present:

def __init__(
self,
model_file: "Path",
output_def: ModelOutputDef,
*args: list,
load_prefix: str = "load",
default_tf_graph: bool = False,
auto_batch_size: bool | int | AutoBatchSize = False,
input_map: dict | None = None,
neighbor_list: "ase.neighborlist.NeighborList | None" = None,
**kwargs: dict,
) -> None:
self.graph = self._load_graph(
model_file,
prefix=load_prefix,
default_tf_graph=default_tf_graph,
input_map=input_map,
)
self.load_prefix = load_prefix
# graph_compatable should be called after graph and prefix are set
if not self._graph_compatable():
raise RuntimeError(
f"model in graph (version {self.model_version}) is incompatible"
f"with the model (version {MODEL_VERSION}) supported by the current code."
"See https://deepmd.rtfd.io/compatibility/ for details."
)
# set default to False, as subclasses may not support
if isinstance(auto_batch_size, bool):
if auto_batch_size:
self.auto_batch_size = AutoBatchSize()
else:
self.auto_batch_size = None
elif isinstance(auto_batch_size, int):
self.auto_batch_size = AutoBatchSize(auto_batch_size)
elif isinstance(auto_batch_size, AutoBatchSize):
self.auto_batch_size = auto_batch_size
else:
raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize")
self.neighbor_list = neighbor_list

# make natoms_vec and default_mesh
if self.neighbor_list is None:
natoms_vec = self.make_natoms_vec(atom_types)
assert natoms_vec[0] == natoms
mesh = make_default_mesh(pbc, not self._check_mixed_types(atom_types))
ghost_map = None
else:
if nframes > 1:
raise NotImplementedError(
"neighbor_list does not support multiple frames"
)
(
natoms_vec,
coords,
atom_types,
mesh,
imap,
ghost_map,
) = self.build_neighbor_list(
coords,
cells if cells is not None else None,
atom_types,
imap,
self.neighbor_list,
)

The lower TF2 model layer also has a neighbor_list argument in call_common, so silently dropping the inference-level argument is especially surprising:

def call_common(
self,
coord: xp.ndarray,
atype: xp.ndarray,
box: xp.ndarray | None = None,
fparam: xp.ndarray | None = None,
aparam: xp.ndarray | None = None,
do_atomic_virial: bool = False,
coord_corr_for_virial: xp.ndarray | None = None,
charge_spin: xp.ndarray | None = None,
neighbor_list: NeighborList | None = None,
) -> dict[str, xp.ndarray]:
return super().call_common(
to_tensorflow_array(coord),
to_tensorflow_array(atype),
box=to_tensorflow_array(box),
fparam=to_tensorflow_array(fparam),
aparam=to_tensorflow_array(aparam),
do_atomic_virial=do_atomic_virial,
coord_corr_for_virial=to_tensorflow_array(coord_corr_for_virial),
charge_spin=to_tensorflow_array(charge_spin),
neighbor_list=neighbor_list,
)

Impact

Users can pass a custom ASE neighbor list to TF2 inference and receive results from the native SavedModel neighbor-list path instead. The argument is a silent no-op, so workflows that rely on externally controlled neighbor construction can be wrong without any warning.

Suggested fix

Either implement neighbor-list forwarding for TF2 inference, or reject non-None neighbor_list in DeepEval.__init__ with a clear NotImplementedError. Add a regression test that passes a sentinel/custom neighbor-list object and verifies it is used or explicitly rejected.

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