Found during a Codex global scan of deepmodeling/deepmd-kit at commit 73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.

Problem

collate_lmdb_frames() takes every find_* flag from the first frame and applies it to the whole batch.

Evidence:

• LmdbDataReader.__getitem__() sets find_<key> per frame when registered data requirements are present or missing:

  deepmd-kit/deepmd/dpmodel/utils/lmdb_data.py

  Lines 648 to 689 in 73de44b

  	# Handle registered data requirements: fill defaults for missing keys,
  	# apply repeat, and cast dtype.
  	for req_key, req_item in self._data_requirements.items():
  	# Extract requirement fields (support both dict and object)
  	if isinstance(req_item, dict):
  	ndof = req_item["ndof"]
  	default = req_item["default"]
  	atomic = req_item["atomic"]
  	repeat = req_item.get("repeat", 1)
  	req_dtype = req_item.get("dtype")
  	if req_dtype is None:
  	req_dtype = (
  	GLOBAL_ENER_FLOAT_PRECISION
  	if req_item.get("high_prec", False)
  	else GLOBAL_NP_FLOAT_PRECISION
  	)
  	else:
  	ndof = req_item.ndof
  	default = req_item.default
  	atomic = req_item.atomic
  	repeat = getattr(req_item, "repeat", 1)
  	req_dtype = req_item.dtype
  	if req_dtype is None:
  	req_dtype = (
  	GLOBAL_ENER_FLOAT_PRECISION
  	if req_item.high_prec
  	else GLOBAL_NP_FLOAT_PRECISION
  	)
  	if req_key not in frame:
  	frame[f"find_{req_key}"] = np.float32(0.0)
  	if atomic:
  	shape = (frame_natoms, ndof)
  	else:
  	shape = (ndof,)
  	data = np.full(shape, default, dtype=req_dtype)
  	if repeat != 1:
  	data = np.repeat(data, repeat).reshape(-1)
  	frame[req_key] = data
  	else:
  	if f"find_{req_key}" not in frame:
  	frame[f"find_{req_key}"] = np.float32(1.0)

• It also sets per-frame find_fparam, find_aparam, find_spin, and find_charge_spin:

  deepmd-kit/deepmd/dpmodel/utils/lmdb_data.py

  Lines 696 to 701 in 73de44b

  	# Add find_* for fparam/aparam/spin/charge_spin if not already set
  	for extra_key in ["fparam", "aparam", "spin", "charge_spin"]:
  	if f"find_{extra_key}" not in frame:
  	frame[f"find_{extra_key}"] = (
  	np.float32(1.0) if extra_key in frame else np.float32(0.0)
  	)

• The collator documents that find_* flags are taken from the first frame:

  deepmd-kit/deepmd/dpmodel/utils/lmdb_data.py

  Lines 811 to 821 in 73de44b

  	def collate_lmdb_frames(frames: list[dict[str, Any]]) -> dict[str, Any]:
  	"""Stack a list of per-frame dicts into a single batch dict.
  	Backend-agnostic via ``array_api_compat``: works for numpy, torch, jax,
  	etc. The array library is inferred from the first frame's ``coord``.
  	Conventions match :func:`deepmd.dpmodel.utils.batch.normalize_batch`:
  	``find_*`` flags are taken from the first frame (constant within a
  	batch); ``fid`` is collected as a list; ``type`` is dropped (callers
  	should already use ``atype``); other arrays are stacked along axis 0.
  	A ``sid`` placeholder is appended.

• The implementation copies frames[0][key] for all find_* keys:

  deepmd-kit/deepmd/dpmodel/utils/lmdb_data.py

  Lines 831 to 841 in 73de44b

  	for key in frames[0]:
  	if key.startswith("find_"):
  	out[key] = frames[0][key]
  	elif key == "fid":
  	out[key] = [f[key] for f in frames]
  	elif key == "type":
  	continue
  	elif frames[0][key] is None:
  	out[key] = None
  	else:
  	out[key] = xp.stack([f[key] for f in frames])

Impact

If optional labels vary within a batch, valid labels in later frames can be ignored, or default-filled labels in later frames can be treated as real data. This can bias losses and validation metrics for partially labeled LMDB datasets.

Suggested Fix

Stack find_* flags per frame, or batch/group frames by label availability before collation. Add a two-frame LMDB batch test where only one frame has a label such as energy or force, and assert the loss mask uses exactly that frame.