feat(lmdb): support mixed-size batches and lazy label availability (#5962)

## Summary

- support `mix:N` LMDB batches containing frames with different atom
counts, using padded rectangular batches for dense models and a flat
ragged node axis for eligible graph models
- compact phantom atoms before graph-model evaluation and make loss
reductions, validation weighting, and epoch sizing use real atom counts
- resolve label availability lazily after data requirements are
registered, so required, optional, defaulted, and partially available
fields are handled without an eager full-dataset scan
- keep non-mixing and native-spin models on their existing rectangular
public paths; the ragged regression coverage uses upstream DPA1 and
avoids model-specific dependencies

## Behavioral changes

- Masked per-atom loss terms now pool all included labels across the
batch instead of averaging per-frame means. This intentionally retires
the bit-identical reduction guarantee from #5738/#5783 for existing
`mixed_type` NPY datasets whose frames have different real atom counts:
those frames are weighted by their real label counts rather than
equally. Uniform-atom-count batches are unchanged. Hessian pair terms
remain normalized per frame so their quadratic component count does not
make large structures dominate a batch.
- Legacy LMDB files have no exact per-frame label-availability metadata.
To avoid an eager O(N) startup scan, the reader uses a bounded probe and
conservatively reduces per-frame `find_*` flags at collation. A missed
rare signature may discard valid supervision for the affected batch, but
default-filled values are never treated as real labels. Recording exact
availability metadata when generating LMDB datasets is tracked in #5954.

## Testing

- all pre-commit hooks passed for the changed files
- 375 passed, 2 skipped, 1 deselected, and 13 subtests passed in the
main targeted LMDB/PT/PT-expt/model suite
- 15 passed in the isolated loss-reduction and decoder-pool regression
suite
- review fixes: 100 passed, 2 skipped, and 2 subtests passed in the
common loss suite; 60 passed in the PT padding-loss suite; all 25 LMDB
training tests passed; padded/unpadded DPA2 graph and Hessian parity
tests passed


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added LMDB batching for frames with different atom counts, including
mixed-size and ragged layouts.
- Added ragged-batch inference and training for supported energy and
spin models.
- Added configurable data-source policies for optional labels and
parameters.
- Added safer handling of padded atoms across neighbor graphs and model
outputs.

- **Bug Fixes**
  - Improved per-atom loss normalization for uneven and padded batches.
- Prevented padded atoms from affecting neighbor searches, metrics, or
losses.
  - Improved handling of missing labels and default-valued data.

- **Documentation**
- Documented mixed-size batching, ragged data, and per-atom
normalization.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->


Closes #5965

Author: OutisLi

deepmd/dpmodel/array_api.py

@@ -11,6 +11,10 @@
     Version,
 )
 
+from deepmd.dpmodel.common import (
+    to_numpy_array,
+)
+
 # Type alias for array_api compatible arrays
 Array = np.ndarray | Any  # Any to support JAX, PyTorch, etc. arrays
 
@@ -27,22 +31,28 @@ def xp_asarray_nodetach(
     ``torch.asarray`` detaches its input from the autograd graph, so calling
     ``xp.asarray`` on a weight attribute that is already a backend tensor
     (e.g. a ``torch.nn.Parameter`` registered by the pt_expt backend)
-    silently breaks gradient flow to that weight.  This helper converts
-    genuine non-backend data (numpy arrays, python scalars/lists) via
-    ``xp.asarray``; backend tensors are returned as-is, with an optional
-    differentiable dtype cast via ``xp.astype``.
-
-    The ``device`` argument only applies to the conversion path: backend
-    tensors are assumed to already live on the working device (they are
-    created together with the inputs).
+    silently breaks gradient flow to that weight. Backend tensors already in
+    ``xp`` are therefore returned as-is, with an optional differentiable dtype
+    cast via ``xp.astype``.
+
+    An array from another namespace cannot retain its autograd graph. It is
+    converted through NumPy before entering ``xp``; this also performs the
+    required device-to-host copy when a CUDA-backed model constant is consumed
+    by a NumPy statistics path.
+
+    The ``device`` argument only applies to the conversion path. Arrays already
+    in ``xp`` are assumed to live on the working device because model buffers
+    and inputs are moved together.
     """
-    if isinstance(obj, np.ndarray) or not array_api_compat.is_array_api_obj(obj):
-        if dtype is None:
-            return xp.asarray(obj, device=device)
-        return xp.asarray(obj, dtype=dtype, device=device)
-    if dtype is not None and obj.dtype != dtype:
-        obj = xp.astype(obj, dtype)
-    return obj
+    if array_api_compat.is_array_api_obj(obj):
+        if array_api_compat.array_namespace(obj) is xp:
+            if dtype is not None and obj.dtype != dtype:
+                obj = xp.astype(obj, dtype)
+            return obj
+        obj = to_numpy_array(obj)
+    if dtype is None:
+        return xp.asarray(obj, device=device)
+    return xp.asarray(obj, dtype=dtype, device=device)
 
 
 # array api adds take_along_axis in https://github.com/data-apis/array-api/pull/816

deepmd/dpmodel/descriptor/dpa4.py

@@ -2186,7 +2186,7 @@ def _canonicalize_charge_spin(
                 raise ValueError("`charge_spin` is required for this SeZM descriptor.")
             charge_spin = xp.reshape(
                 xp_asarray_nodetach(
-                    xp, np.asarray(self.default_chg_spin), dtype=dtype, device=device
+                    xp, self.default_chg_spin, dtype=dtype, device=device
                 ),
                 (1, 2),
             )

deepmd/dpmodel/loss/dos.py

@@ -159,7 +159,7 @@ def call(
             )
             diff3d = local_pred - local_label  # [nf, natoms, numb_dos]
             if "mask" in model_dict:
-                # idiom 1: per-frame masked mean, then average over frames
+                # Idiom 1 (per-atom masked mean, ncomp=numb_dos).
                 maskf = xp.astype(model_dict["mask"], diff3d.dtype)  # [nf, natoms]
                 l2_local_loss_dos = masked_atom_mean(
                     xp.square(diff3d), maskf, self.numb_dos
@@ -184,7 +184,7 @@ def call(
             )
             diff3d = local_pred_cdf - local_label_cdf  # [nf, natoms, numb_dos]
             if "mask" in model_dict:
-                # idiom 1: per-frame masked mean, then average over frames
+                # Idiom 1 (per-atom masked mean, ncomp=numb_dos).
                 maskf = xp.astype(model_dict["mask"], diff3d.dtype)  # [nf, natoms]
                 l2_local_loss_cdf = masked_atom_mean(
                     xp.square(diff3d), maskf, self.numb_dos