fix bugs

Author: Han Wang

deepmd/jax/infer/deep_eval.py

@@ -388,7 +388,9 @@ def _eval_model(
 
         results = []
         for odef in request_defs:
-            dp_name = self._OUTDEF_DP2BACKEND[odef.name]
+            # HLO and TFModelWrapper return raw internal keys (not translated),
+            # so no key mapping is needed here.
+            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:

source/tests/universal/dpmodel/backend.py

@@ -21,6 +21,8 @@ class DPTestCase(BackendTestCase):
     """DP module to test."""
 
     def forward_wrapper(self, x):
+        if not hasattr(x, "forward_lower") and hasattr(x, "call_lower"):
+            x.forward_lower = x.call_lower
         return x
 
     def forward_wrapper_cpu_ref(self, x):

source/tests/universal/dpmodel/model/test_model.py

@@ -164,7 +164,7 @@ def setUpClass(cls) -> None:
             ft,
             type_map=cls.expected_type_map,
         )
-        cls.output_def = cls.module.model_output_def().get_data()
+        cls.output_def = cls.module.translated_output_def()
         cls.expected_has_message_passing = ds.has_message_passing()
         cls.expected_sel_type = ft.get_sel_type()
         cls.expected_dim_fparam = ft.get_dim_fparam()