fix(tf): fix serialization of dipole fitting with sel_type

.gitignore

@@ -55,3 +55,12 @@ node_modules/
 test_dp_test/
 test_dp_test_*.out
 *_detail.out
+
+# Training and model output files
+*.pth
+*.ckpt*
+checkpoint
+lcurve.out
+out.json
+input_v2_compat.json
+frozen_model.*

deepmd/dpmodel/utils/network.py

@@ -935,7 +935,7 @@ def __init__(
         self._networks = [None for ii in range(ntypes**ndim)]
         for ii, network in enumerate(networks):
             self[ii] = network
-        if len(networks):
+        if len(networks) and all(net is not None for net in networks):
             self.check_completeness()
 
     def check_completeness(self) -> None:
@@ -969,7 +969,9 @@ def __getitem__(self, key):
         return self._networks[self._convert_key(key)]
 
     def __setitem__(self, key, value) -> None:
-        if isinstance(value, self.network_type):
+        if value is None:
+            pass
+        elif isinstance(value, self.network_type):
             pass
         elif isinstance(value, dict):
             value = self.network_type.deserialize(value)
@@ -993,7 +995,9 @@ def serialize(self) -> dict:
             "ndim": self.ndim,
             "ntypes": self.ntypes,
             "network_type": network_type_name,
-            "networks": [nn.serialize() for nn in self._networks],
+            "networks": [
+                nn.serialize() if nn is not None else None for nn in self._networks
+            ],
         }
 
     @classmethod

deepmd/tf/fit/dipole.py

@@ -422,6 +422,7 @@ def serialize(self, suffix: str) -> dict:
             "activation_function": self.activation_function_name,
             "precision": self.fitting_precision.name,
             "exclude_types": [],
+            "sel_type": self.sel_type,
             "nets": self.serialize_network(
                 ntypes=self.ntypes,
                 ndim=0 if self.mixed_types else 1,
@@ -434,6 +435,13 @@ def serialize(self, suffix: str) -> dict:
                 trainable=self.trainable,
                 suffix=suffix,
             ),
+            "@variables": {
+                "fparam_avg": self.fparam_avg,
+                "fparam_inv_std": self.fparam_inv_std,
+                "aparam_avg": self.aparam_avg,
+                "aparam_inv_std": self.aparam_inv_std,
+                "case_embd": None,
+            },
             "type_map": self.type_map,
         }
         return data

deepmd/tf/fit/fitting.py

@@ -244,7 +244,9 @@ def deserialize_network(cls, data: dict, suffix: str = "") -> dict:
             else:
                 raise ValueError(f"Invalid ndim: {fittings.ndim}")
             network = fittings[net_idx]
-            assert network is not None
+            if network is None:
+                # Skip types that are not selected (when sel_type is used)
+                continue
             for layer_idx, layer in enumerate(network.layers):
                 if layer_idx == len(network.layers) - 1:
                     layer_name = "final_layer"