fixes

Author: Han Wang

deepmd/jax/infer/deep_eval.py

@@ -388,7 +388,7 @@ def _eval_model(
 
         results = []
         for odef in request_defs:
-            dp_name = self._OUTDEF_DP2BACKEND[odef.name]
+            dp_name = odef.name
             if dp_name in batch_output:
                 shape = self._get_output_shape(odef, nframes, natoms)
                 if batch_output[dp_name] is not None:

deepmd/jax/jax2tf/serialization.py

@@ -34,7 +34,7 @@ def deserialize_to_file(model_file: str, data: dict) -> None:
     if model_file.endswith(".savedmodel"):
         model = BaseModel.deserialize(data["model"])
         model_def_script = data["model_def_script"]
-        call_lower = model.call_lower
+        call_lower = model.call_common_lower
 
         tf_model = tf.Module()
 

deepmd/jax/utils/serialization.py

@@ -49,7 +49,7 @@ def deserialize_to_file(model_file: str, data: dict) -> None:
     elif model_file.endswith(".hlo"):
         model = BaseModel.deserialize(data["model"])
         model_def_script = data["model_def_script"]
-        call_lower = model.call_lower
+        call_lower = model.call_common_lower
 
         nf, nloc, nghost = jax_export.symbolic_shape("nf, nloc, nghost")
 

source/tests/common/dpmodel/test_padding_atoms.py

@@ -69,8 +69,8 @@ def test_padding_atoms_consistency(self):
         result = model.call(*args)
         # test intensive
         np.testing.assert_allclose(
-            result[f"{var_name}_redu"],
-            np.mean(result[f"{var_name}"], axis=1),
+            result[var_name],
+            np.mean(result[f"atom_{var_name}"], axis=1),
             atol=self.atol,
         )
         # test padding atoms
@@ -93,8 +93,8 @@ def test_padding_atoms_consistency(self):
             args = [coord_padding, atype_padding, self.cell]
             result_padding = model.call(*args)
             np.testing.assert_allclose(
-                result[f"{var_name}_redu"],
-                result_padding[f"{var_name}_redu"],
+                result[var_name],
+                result_padding[var_name],
                 atol=self.atol,
             )
 

source/tests/consistent/model/test_property.py

@@ -186,8 +186,8 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]:
         property_name = self.data["fitting_net"]["property_name"]
         if backend in {self.RefBackend.DP, self.RefBackend.JAX}:
             return (
-                ret[f"{property_name}_redu"].ravel(),
                 ret[property_name].ravel(),
+                ret[f"atom_{property_name}"].ravel(),
             )
         elif backend is self.RefBackend.PT:
             return (

source/tests/jax/test_padding_atoms.py

@@ -89,8 +89,8 @@ def test_padding_atoms_consistency(self):
         result = model.call(*args)
         # test intensive
         np.testing.assert_allclose(
-            to_numpy_array(result[f"{var_name}_redu"]),
-            np.mean(to_numpy_array(result[f"{var_name}"]), axis=1),
+            to_numpy_array(result[var_name]),
+            np.mean(to_numpy_array(result[f"atom_{var_name}"]), axis=1),
             atol=self.atol,
         )
         # test padding atoms
@@ -115,8 +115,8 @@ def test_padding_atoms_consistency(self):
             ]
             result_padding = model.call(*args)
             np.testing.assert_allclose(
-                to_numpy_array(result[f"{var_name}_redu"]),
-                to_numpy_array(result_padding[f"{var_name}_redu"]),
+                to_numpy_array(result[var_name]),
+                to_numpy_array(result_padding[var_name]),
                 atol=self.atol,
             )