fix(dpmodel): support parameter shorthand in DeepEval

deepmd/dpmodel/infer/deep_eval.py

@@ -40,6 +40,7 @@
 from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper
 from deepmd.infer.deep_eval import (
     DeepEvalBackend,
+    _standardize_fparam_aparam,
 )
 from deepmd.infer.deep_polar import (
     DeepPolar,
@@ -240,6 +241,14 @@ def eval(
         natoms, numb_test = self._get_natoms_and_nframes(
             coords, atom_types, len(atom_types.shape) > 1
         )
+        fparam, aparam = _standardize_fparam_aparam(
+            fparam,
+            aparam,
+            numb_test,
+            natoms,
+            self.get_dim_fparam(),
+            self.get_dim_aparam(),
+        )
         request_defs = self._get_request_defs(atomic)
         out = self._eval_func(self._eval_model, numb_test, natoms)(
             coords, cells, atom_types, fparam, aparam, request_defs

deepmd/infer/deep_eval.py

@@ -30,6 +30,75 @@
     import ase.neighborlist
 
 
+def _standardize_fparam_aparam(
+    fparam: np.ndarray | list | None,
+    aparam: np.ndarray | list | None,
+    nframes: int,
+    natoms: int,
+    dim_fparam: int,
+    dim_aparam: int,
+) -> tuple[np.ndarray | None, np.ndarray | None]:
+    """Normalize documented parameter shorthand to frame-major arrays.
+
+    This normalization must happen before automatic batching. In particular,
+    an ``(natoms, dim_aparam)`` shared atomic parameter has an atom axis first;
+    a batcher would otherwise mistake that axis for frames and slice it.
+
+    Parameters
+    ----------
+    fparam : np.ndarray or list or None
+        Frame parameters in full ``(nframes, dim_fparam)`` form or shared
+        ``(dim_fparam,)`` shorthand.
+    aparam : np.ndarray or list or None
+        Atomic parameters in full ``(nframes, natoms, dim_aparam)`` form,
+        shared-per-atom ``(natoms, dim_aparam)`` shorthand, or shared-for-all
+        ``(dim_aparam,)`` shorthand.
+    nframes : int
+        Number of input frames.
+    natoms : int
+        Number of atoms in each frame.
+    dim_fparam : int
+        Number of frame-parameter components.
+    dim_aparam : int
+        Number of atomic-parameter components.
+
+    Returns
+    -------
+    normalized_fparam : np.ndarray or None
+        Frame parameters with shape ``(nframes, dim_fparam)`` when provided.
+    normalized_aparam : np.ndarray or None
+        Atomic parameters with shape ``(nframes, natoms, dim_aparam)`` when
+        provided. The public wrapper may subsequently flatten the last two
+        axes to preserve its historical backend ABI.
+    """
+    if fparam is not None:
+        fparam = np.asarray(fparam)
+        if fparam.size == nframes * dim_fparam:
+            fparam = fparam.reshape(nframes, dim_fparam)
+        elif fparam.size == dim_fparam:
+            fparam = np.tile(fparam.reshape(1, dim_fparam), (nframes, 1))
+        else:
+            raise RuntimeError(
+                "got wrong size of frame param, should be either "
+                f"{nframes} x {dim_fparam} or {dim_fparam}"
+            )
+    if aparam is not None:
+        aparam = np.asarray(aparam)
+        if aparam.size == nframes * natoms * dim_aparam:
+            aparam = aparam.reshape(nframes, natoms, dim_aparam)
+        elif aparam.size == natoms * dim_aparam:
+            aparam = np.tile(aparam.reshape(1, natoms, dim_aparam), (nframes, 1, 1))
+        elif aparam.size == dim_aparam:
+            aparam = np.tile(aparam.reshape(1, 1, dim_aparam), (nframes, natoms, 1))
+        else:
+            raise RuntimeError(
+                "got wrong size of atomic param, should be either "
+                f"{nframes} x {natoms} x {dim_aparam} or "
+                f"{natoms} x {dim_aparam} or {dim_aparam}"
+            )
+    return fparam, aparam
+
+
 class DeepEvalBackend(ABC):
     """Low-level Deep Evaluator interface.
 
@@ -948,28 +1017,18 @@ def _standard_input(
         coords = coords.reshape(nframes, natoms, 3)
         if cells is not None:
             cells = cells.reshape(nframes, 3, 3)
-        if fparam is not None:
-            fdim = self.get_dim_fparam()
-            if fparam.size == nframes * fdim:
-                fparam = np.reshape(fparam, [nframes, fdim])
-            elif fparam.size == fdim:
-                fparam = np.tile(fparam.reshape([-1]), [nframes, 1])
-            else:
-                raise RuntimeError(
-                    f"got wrong size of frame param, should be either {nframes} x {fdim} or {fdim}"
-                )
+        fparam, aparam = _standardize_fparam_aparam(
+            fparam,
+            aparam,
+            nframes,
+            natoms,
+            self.get_dim_fparam(),
+            self.get_dim_aparam(),
+        )
         if aparam is not None:
-            fdim = self.get_dim_aparam()
-            if aparam.size == nframes * natoms * fdim:
-                aparam = np.reshape(aparam, [nframes, natoms * fdim])
-            elif aparam.size == natoms * fdim:
-                aparam = np.tile(aparam.reshape([-1]), [nframes, 1])
-            elif aparam.size == fdim:
-                aparam = np.tile(aparam.reshape([-1]), [nframes, natoms])
-            else:
-                raise RuntimeError(
-                    f"got wrong size of frame param, should be either {nframes} x {natoms} x {fdim} or {natoms} x {fdim} or {fdim}"
-                )
+            # Preserve the historical flattened backend ABI used by the public
+            # wrapper; backend adapters normalize it back to frame-major 3-D.
+            aparam = aparam.reshape(nframes, natoms * self.get_dim_aparam())
         return coords, cells, atom_types, fparam, aparam, nframes, natoms
 
     def get_sel_type(self) -> list[int]:

deepmd/jax/infer/deep_eval.py

@@ -35,6 +35,7 @@
 from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper
 from deepmd.infer.deep_eval import (
     DeepEvalBackend,
+    _standardize_fparam_aparam,
 )
 from deepmd.infer.deep_polar import (
     DeepPolar,
@@ -278,6 +279,14 @@ def eval(
         natoms, numb_test = self._get_natoms_and_nframes(
             coords, atom_types, len(atom_types.shape) > 1
         )
+        fparam, aparam = _standardize_fparam_aparam(
+            fparam,
+            aparam,
+            numb_test,
+            natoms,
+            self.get_dim_fparam(),
+            self.get_dim_aparam(),
+        )
         request_defs = self._get_request_defs(atomic)
         out = self._eval_func(self._eval_model, numb_test, natoms)(
             coords, cells, atom_types, fparam, aparam, charge_spin, request_defs

deepmd/pd/infer/deep_eval.py

@@ -28,6 +28,7 @@
 from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper
 from deepmd.infer.deep_eval import (
     DeepEvalBackend,
+    _standardize_fparam_aparam,
 )
 from deepmd.infer.deep_polar import (
     DeepGlobalPolar,
@@ -376,6 +377,14 @@ def eval(
         natoms, numb_test = self._get_natoms_and_nframes(
             coords, atom_types, len(atom_types.shape) > 1
         )
+        fparam, aparam = _standardize_fparam_aparam(
+            fparam,
+            aparam,
+            numb_test,
+            natoms,
+            self.get_dim_fparam(),
+            self.get_dim_aparam(),
+        )
         request_defs = self._get_request_defs(atomic)
         if "spin" not in kwargs or kwargs["spin"] is None:
             out = self._eval_func(self._eval_model, numb_test, natoms)(

deepmd/pt/infer/deep_eval.py

@@ -29,6 +29,7 @@
 from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper
 from deepmd.infer.deep_eval import (
     DeepEvalBackend,
+    _standardize_fparam_aparam,
 )
 from deepmd.infer.deep_polar import (
     DeepGlobalPolar,
@@ -544,6 +545,14 @@ def eval(
         natoms, numb_test = self._get_natoms_and_nframes(
             coords, atom_types, len(atom_types.shape) > 1
         )
+        fparam, aparam = _standardize_fparam_aparam(
+            fparam,
+            aparam,
+            numb_test,
+            natoms,
+            self.get_dim_fparam(),
+            self.get_dim_aparam(),
+        )
         request_defs = self._get_request_defs(atomic)
         if "spin" not in kwargs or kwargs["spin"] is None:
             out = self._eval_func(self._eval_model, numb_test, natoms)(
@@ -1312,6 +1321,17 @@ def eval_embedding(
         natoms, numb_test = self._get_natoms_and_nframes(
             coords, atom_types, len(atom_types.shape) > 1
         )
+        # Normalize shared parameter shorthand before auto batching. Otherwise
+        # a one-dimensional fparam/aparam is passed unchanged to every split,
+        # and _eval_embedding cannot reshape it to the split frame count.
+        fparam, aparam = _standardize_fparam_aparam(
+            fparam,
+            aparam,
+            numb_test,
+            natoms,
+            self.get_dim_fparam(),
+            self.get_dim_aparam(),
+        )
         return self._eval_func(self._eval_embedding, numb_test, natoms)(
             coords, cells, atom_types, fparam, aparam, charge_spin, dtype
         )

deepmd/pt_expt/infer/deep_eval.py

@@ -41,6 +41,7 @@
 from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper
 from deepmd.infer.deep_eval import (
     DeepEvalBackend,
+    _standardize_fparam_aparam,
 )
 from deepmd.infer.deep_polar import (
     DeepPolar,
@@ -848,10 +849,15 @@ def eval(
             Calculate the atomic energy and virial
         fparam
             The frame parameter.
-            The array should be of size nframes x dim_fparam.
+            The array can be of size :
+            - nframes x dim_fparam.
+            - dim_fparam. Then all frames are assumed to be provided with the same fparam.
         aparam
-            The atomic parameter.
-            The array should be of size nframes x natoms x dim_aparam.
+            The atomic parameter
+            The array can be of size :
+            - nframes x natoms x dim_aparam.
+            - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam.
+            - dim_aparam. Then all frames and atoms are provided with the same aparam.
         charge_spin
             The charge and spin values for each frame.
             The array should be reshape-compatible with nframes x 2, where the first
@@ -875,6 +881,14 @@ def eval(
         natoms, numb_test = self._get_natoms_and_nframes(
             coords, atom_types, len(atom_types.shape) > 1
         )
+        fparam, aparam = _standardize_fparam_aparam(
+            fparam,
+            aparam,
+            numb_test,
+            natoms,
+            self.get_dim_fparam(),
+            self.get_dim_aparam(),
+        )
         request_defs = self._get_request_defs(atomic)
         spins = kwargs.get("spin")
         if self._is_spin and spins is None:
@@ -2203,6 +2217,19 @@ def eval_descriptor(
                 "eval_descriptor is not supported for this model type "
                 f"({type(self._dpmodel).__name__})."
             )
+        nframes = coords.shape[0]
+        natoms = len(atom_types) if len(atom_types.shape) == 1 else atom_types.shape[1]
+        # Canonicalize shared parameter shorthands before _prepare_nlist_inputs
+        # reshapes them; otherwise a shared per-atom array can be mistaken for
+        # a batch of frames.
+        fparam, aparam = _standardize_fparam_aparam(
+            fparam,
+            aparam,
+            nframes,
+            natoms,
+            self.get_dim_fparam(),
+            self.get_dim_aparam(),
+        )
         (
             ext_coord_t,
             ext_atype_t,
@@ -2274,6 +2301,19 @@ def eval_fitting_last_layer(
                 "eval_fitting_last_layer is not supported for this model type "
                 f"({type(self._dpmodel).__name__})."
             )
+        nframes = coords.shape[0]
+        natoms = len(atom_types) if len(atom_types.shape) == 1 else atom_types.shape[1]
+        # Canonicalize shared parameter shorthands before _prepare_nlist_inputs
+        # reshapes them; otherwise a shared per-atom array can be mistaken for
+        # a batch of frames.
+        fparam, aparam = _standardize_fparam_aparam(
+            fparam,
+            aparam,
+            nframes,
+            natoms,
+            self.get_dim_fparam(),
+            self.get_dim_aparam(),
+        )
         (
             ext_coord_t,
             ext_atype_t,

deepmd/tf/infer/deep_eval.py

@@ -32,6 +32,7 @@
 )
 from deepmd.infer.deep_eval import (
     DeepEvalBackend,
+    _standardize_fparam_aparam,
 )
 from deepmd.infer.deep_polar import (
     DeepGlobalPolar,
@@ -761,6 +762,14 @@ def eval(
             coords,
             atom_types,
         )
+        fparam, aparam = _standardize_fparam_aparam(
+            fparam,
+            aparam,
+            numb_test,
+            natoms,
+            self.get_dim_fparam(),
+            self.get_dim_aparam(),
+        )
         output = self._eval_func(self._eval_inner, numb_test, natoms)(
             coords,
             cells,
@@ -1084,6 +1093,17 @@ def eval_descriptor(
             coords,
             atom_types,
         )
+        # Canonicalize shared parameter shorthands before AutoBatchSize slices
+        # the frame axis; otherwise a shared per-atom array can be mistaken for
+        # a batch of frames.
+        fparam, aparam = _standardize_fparam_aparam(
+            fparam,
+            aparam,
+            numb_test,
+            natoms,
+            self.get_dim_fparam(),
+            self.get_dim_aparam(),
+        )
         descriptor = self._eval_func(self._eval_descriptor_inner, numb_test, natoms)(
             coords,
             cells,

deepmd/tf2/infer/deep_eval.py

@@ -32,6 +32,7 @@
 from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper
 from deepmd.infer.deep_eval import (
     DeepEvalBackend,
+    _standardize_fparam_aparam,
 )
 from deepmd.infer.deep_polar import (
     DeepPolar,
@@ -292,6 +293,14 @@ def eval(
         natoms, numb_test = self._get_natoms_and_nframes(
             coords, atom_types, len(atom_types.shape) > 1
         )
+        fparam, aparam = _standardize_fparam_aparam(
+            fparam,
+            aparam,
+            numb_test,
+            natoms,
+            self.get_dim_fparam(),
+            self.get_dim_aparam(),
+        )
         request_defs = self._get_request_defs(atomic)
         out = self._eval_func(self._eval_model, numb_test, natoms)(
             coords, cells, atom_types, fparam, aparam, request_defs