fix uts

Author: iProzd

deepmd/dpmodel/fitting/general_fitting.py

@@ -292,6 +292,8 @@ def __setitem__(self, key, value) -> None:
             self.case_embd = value
         elif key in ["scale"]:
             self.scale = value
+        elif key in ["default_fparam_tensor"]:
+            self.default_fparam_tensor = value
         else:
             raise KeyError(key)
 
@@ -310,6 +312,8 @@ def __getitem__(self, key):
             return self.case_embd
         elif key in ["scale"]:
             return self.scale
+        elif key in ["default_fparam_tensor"]:
+            return self.default_fparam_tensor
         else:
             raise KeyError(key)
 

deepmd/jax/fitting/fitting.py

@@ -35,6 +35,7 @@ def setattr_for_general_fitting(name: str, value: Any) -> Any:
         "fparam_inv_std",
         "aparam_avg",
         "aparam_inv_std",
+        "default_fparam_tensor",
     }:
         value = to_jax_array(value)
         if value is not None:

deepmd/pt/model/task/fitting.py

@@ -495,6 +495,8 @@ def __setitem__(self, key, value) -> None:
             self.case_embd = value
         elif key in ["scale"]:
             self.scale = value
+        elif key in ["default_fparam_tensor"]:
+            self.default_fparam_tensor = value
         else:
             raise KeyError(key)
 
@@ -513,6 +515,8 @@ def __getitem__(self, key):
             return self.case_embd
         elif key in ["scale"]:
             return self.scale
+        elif key in ["default_fparam_tensor"]:
+            return self.default_fparam_tensor
         else:
             raise KeyError(key)
 

source/tests/consistent/fitting/test_ener.py

@@ -138,6 +138,18 @@ def skip_pd(self) -> bool:
         # so skip this in CI test
         return not INSTALLED_PD or precision == "bfloat16" or default_fparam is not None
 
+    @property
+    def skip_tf(self) -> bool:
+        (
+            resnet_dt,
+            precision,
+            mixed_types,
+            (numb_fparam, default_fparam),
+            (numb_aparam, use_aparam_as_mask),
+            atom_ener,
+        ) = self.param
+        return not INSTALLED_TF or default_fparam is not None
+
     tf_class = EnerFittingTF
     dp_class = EnerFittingDP
     pt_class = EnerFittingPT