fix(tf): throw RuntimeError for se_a + type_embedding

deepmd/tf/descriptor/se_a.py

@@ -301,6 +301,7 @@
                 self.stat_descrpt *= tf.reshape(mask, tf.shape(self.stat_descrpt))
         self.sub_sess = tf.Session(graph=sub_graph, config=default_tf_session_config)
         self.original_sel = None
+        self.use_tebd: Optional[bool] = None
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
@@ -746,8 +747,10 @@
     ):
         if input_dict is not None:
             type_embedding = input_dict.get("type_embedding", None)
+            self.use_tebd = True
         else:
             type_embedding = None
+            self.use_tebd = False
         if self.stripped_type_embedding and type_embedding is None:
             raise RuntimeError("type_embedding is required for se_a_tebd_v2 model.")
         start_index = 0
@@ -1406,7 +1409,7 @@
             raise NotImplementedError(
                 "Serialization is unsupported when tebd_input_mode is set to 'strip'"
             )
-        if (self.original_sel != self.sel_a).any():
+        if self.original_sel is not None and (self.original_sel != self.sel_a).any():
             raise NotImplementedError(
                 "Adjusting sel is unsupported by the native model"
             )
@@ -1416,9 +1419,11 @@
             raise NotImplementedError("spin is unsupported")
         assert self.davg is not None
         assert self.dstd is not None
-        # TODO: tf: handle type embedding in DescrptSeA.serialize
-        # not sure how to handle type embedding - type embedding is not a model parameter,
-        # but instead a part of the input data. Maybe the interface should be refactored...
+        assert self.use_tebd is not None
+        if self.use_tebd:
+            raise RuntimeError(
+                "Serialization is unsupported when type_embedding is used."
+            )
 
         return {
             "@class": "Descriptor",

source/tests/tf/test_model_se_a_type.py

@@ -179,3 +179,6 @@ def test_model(self):
         np.testing.assert_almost_equal(e, refe, places)
         np.testing.assert_almost_equal(f, reff, places)
         np.testing.assert_almost_equal(v, refv, places)
+
+        with self.assertRaises(RuntimeError):
+            descrpt.serialize(suffix="se_a_type")