fix: UTs

anyangml · Feb 22, 2024 · 6052f7e · 6052f7e
1 parent 13df1db
commit 6052f7e
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 6 deletions.
diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py
@@ -344,7 +344,7 @@ def serialize(self, suffix: str) -> dict:
             The serialized data
         """
         data = {
-            "var_name": "energy",
+            "var_name": "dipole",
             "ntypes": self.ntypes,
             "dim_descrpt": self.dim_descrpt,
             "embedding_width": self.dim_rot_mat_1,
@@ -362,6 +362,7 @@ def serialize(self, suffix: str) -> dict:
                 # TODO: consider type embeddings
                 ndim=1,
                 in_dim=self.dim_descrpt,
+                out_dim = self.dim_rot_mat_1,
                 neuron=self.n_neuron,
                 activation_function=self.activation_function_name,
                 resnet_dt=self.resnet_dt,

diff --git a/deepmd/tf/fit/fitting.py b/deepmd/tf/fit/fitting.py
@@ -7,6 +7,7 @@
     Callable,
     List,
     Type,
+    Optional,
 )
 
 from deepmd.dpmodel.utils.network import (
@@ -175,6 +176,7 @@ def serialize_network(
         activation_function: str,
         resnet_dt: bool,
         variables: dict,
+        out_dim: Optional[int] = 1,
         suffix: str = "",
     ) -> dict:
         """Serialize network.
@@ -197,6 +199,8 @@ def serialize_network(
             The input variables
         suffix : str, optional
             The suffix of the scope
+        out_dim : int, optional
+            The output dimension
 
         Returns
         -------
@@ -231,7 +235,7 @@ def serialize_network(
                 # initialize the network if it is not initialized
                 fittings[network_idx] = FittingNet(
                     in_dim=in_dim,
-                    out_dim=1,
+                    out_dim=out_dim,
                     neuron=neuron,
                     activation_function=activation_function,
                     resnet_dt=resnet_dt,

diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py
@@ -252,8 +252,12 @@ def test_tf_consistent_with_ref(self):
         tf_obj = self.tf_class.deserialize(data1, suffix=self.unique_id)
         ret2, data2 = self.get_tf_ret_serialization_from_cls(tf_obj)
         ret2 = self.extract_ret(ret2, self.RefBackend.TF)
-        if not tf_obj.__class__.__name__.startswith("Polar"):
-            np.testing.assert_equal(data1, data2)  # tf, pt serialization mismatch
+        if tf_obj.__class__.__name__.startswith("Polar", "Dipole"):
+            # tf, pt serialization mismatch
+            common_keys = set(data1.keys()) & set(data2.keys())
+            data1 = {k: data1[k] for k in common_keys}
+            data2 = {k: data2[k] for k in common_keys}
+        np.testing.assert_equal(data1, data2)  
         for rr1, rr2 in zip(ret1, ret2):
             np.testing.assert_allclose(
                 rr1.ravel(), rr2.ravel(), rtol=self.rtol, atol=self.atol
@@ -318,8 +322,12 @@ def test_pt_consistent_with_ref(self):
         ret2 = self.eval_pt(obj)
         ret2 = self.extract_ret(ret2, self.RefBackend.PT)
         data2 = obj.serialize()
-        if not obj.__class__.__name__.startswith("Polar"):
-            np.testing.assert_equal(data1, data2)  # tf, pt serialization mismatch
+        if obj.__class__.__name__.startswith("Polar", "Dipole"):
+            # tf, pt serialization mismatch
+            common_keys = set(data1.keys()) & set(data2.keys())
+            data1 = {k: data1[k] for k in common_keys}
+            data2 = {k: data2[k] for k in common_keys}
+        np.testing.assert_equal(data1, data2)  
         for rr1, rr2 in zip(ret1, ret2):
             np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol)