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

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 @@
     ``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