fix(jax): support charge-spin savedmodel export

deepmd/jax/infer/deep_eval.py

@@ -250,7 +250,11 @@ def eval(
             - 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.
         **kwargs
-            Other parameters
+            Other parameters.
+            charge_spin : array-like, optional
+                The per-frame charge/spin conditioning input. The array can be
+                of size nframes x dim_chg_spin, or dim_chg_spin to reuse the
+                same value for all frames.
 
         Returns
         -------
@@ -259,6 +263,7 @@ def eval(
             variables, and the values are the corresponding output arrays.
         """
         # convert all of the input to numpy array
+        charge_spin = kwargs.pop("charge_spin", None)
         atom_types = np.array(atom_types, dtype=np.int32)
         coords = np.array(coords)
         if cells is not None:
@@ -268,7 +273,7 @@ def eval(
         )
         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
+            coords, cells, atom_types, fparam, aparam, charge_spin, request_defs
         )
         return dict(
             zip(
@@ -361,6 +366,7 @@ def _eval_model(
         atom_types: np.ndarray,
         fparam: np.ndarray | None,
         aparam: np.ndarray | None,
+        charge_spin: np.ndarray | None,
         request_defs: list[OutputVariableDef],
     ) -> tuple[np.ndarray, ...]:
         model = self.dp
@@ -394,17 +400,32 @@ def _eval_model(
             aparam_input = aparam.reshape(nframes, natoms, self.get_dim_aparam())
         else:
             aparam_input = None
+        if charge_spin is not None and not self.has_chg_spin_ebd():
+            raise ValueError(
+                "charge_spin was provided, but this model does not support "
+                "charge/spin conditioning."
+            )
+        charge_spin_input = (
+            np.asarray(charge_spin, dtype=GLOBAL_NP_FLOAT_PRECISION)
+            if self.has_chg_spin_ebd() and charge_spin is not None
+            else None
+        )
 
         do_atomic_virial = any(
             x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs
         )
+        model_kwargs = {
+            "box": to_jax_array(box_input),
+            "fparam": to_jax_array(fparam_input),
+            "aparam": to_jax_array(aparam_input),
+            "do_atomic_virial": do_atomic_virial,
+        }
+        if self.has_chg_spin_ebd():
+            model_kwargs["charge_spin"] = to_jax_array(charge_spin_input)
         batch_output = model(
             to_jax_array(coord_input),
             to_jax_array(type_input),
-            box=to_jax_array(box_input),
-            fparam=to_jax_array(fparam_input),
-            aparam=to_jax_array(aparam_input),
-            do_atomic_virial=do_atomic_virial,
+            **model_kwargs,
         )
         if isinstance(batch_output, tuple):
             batch_output = batch_output[0]
@@ -471,3 +492,21 @@ def get_model(self) -> Any:
     def has_default_fparam(self) -> bool:
         """Check if the model has default frame parameters."""
         return self.dp.has_default_fparam()
+
+    def has_chg_spin_ebd(self) -> bool:
+        """Check if the model has charge spin embedding."""
+        if hasattr(self.dp, "has_chg_spin_ebd"):
+            return self.dp.has_chg_spin_ebd()
+        return False
+
+    def get_dim_chg_spin(self) -> int:
+        """Get the dimension of charge_spin input."""
+        if hasattr(self.dp, "get_dim_chg_spin"):
+            return self.dp.get_dim_chg_spin()
+        return 0
+
+    def has_default_chg_spin(self) -> bool:
+        """Check if the model has default charge_spin values."""
+        if hasattr(self.dp, "has_default_chg_spin"):
+            return self.dp.has_default_chg_spin()
+        return False

deepmd/jax/jax2tf/make_model.py

@@ -28,20 +28,36 @@
     communicate_extended_output,
 )
 
-
-def model_call_from_call_lower(
-    *,  # enforce keyword-only arguments
-    call_lower: Callable[
+CallLower = (
+    Callable[
         [
             tf.Tensor,
             tf.Tensor,
             tf.Tensor,
             tf.Tensor,
             tf.Tensor,
-            bool,
+            tf.Tensor,
         ],
         dict[str, tf.Tensor],
-    ],
+    ]
+    | Callable[
+        [
+            tf.Tensor,
+            tf.Tensor,
+            tf.Tensor,
+            tf.Tensor,
+            tf.Tensor,
+            tf.Tensor,
+            tf.Tensor,
+        ],
+        dict[str, tf.Tensor],
+    ]
+)
+
+
+def model_call_from_call_lower(
+    *,  # enforce keyword-only arguments
+    call_lower: CallLower,
     rcut: float,
     sel: list[int],
     mixed_types: bool,
@@ -51,6 +67,7 @@ def model_call_from_call_lower(
     box: tf.Tensor,
     fparam: tf.Tensor,
     aparam: tf.Tensor,
+    charge_spin: tf.Tensor | None = None,
     do_atomic_virial: bool = False,
 ) -> dict[str, tf.Tensor]:
     """Return model prediction from lower interface."""
@@ -79,13 +96,18 @@ def model_call_from_call_lower(
         distinguish_types=False,
     )
     extended_coord = tf.reshape(extended_coord, [nframes, -1, 3])
+    call_lower_kwargs = {
+        "fparam": fp,
+        "aparam": ap,
+    }
+    if charge_spin is not None:
+        call_lower_kwargs["charge_spin"] = charge_spin
     model_predict_lower = call_lower(
         extended_coord,
         extended_atype,
         nlist,
         mapping,
-        fparam=fp,
-        aparam=ap,
+        **call_lower_kwargs,
     )
     model_predict = communicate_extended_output(
         model_predict_lower,