fix(tf): make dipole, polar, and dos models consistent with dpmodel

deepmd/tf/model/dos.py

@@ -149,6 +149,9 @@ def build(
             t_ver = tf.constant(MODEL_VERSION, name="model_version", dtype=tf.string)
             t_od = tf.constant(self.numb_dos, name="output_dim", dtype=tf.int32)
 
+            # Initialize out_bias and out_std for DOS models
+            self.init_out_stat(suffix=suffix)
+
         coord = tf.reshape(coord_, [-1, natoms[1] * 3])
         atype = tf.reshape(atype_, [-1, natoms[1]])
         input_dict["nframes"] = tf.shape(coord)[0]
@@ -181,6 +184,10 @@ def build(
         atom_dos = self.fitting.build(
             dout, natoms, input_dict, reuse=reuse, suffix=suffix
         )
+
+        # Apply out_bias and out_std directly to DOS output
+        atom_dos = self._apply_out_bias_std(atom_dos, atype, natoms, coord)
+
         self.atom_dos = atom_dos
 
         dos_raw = atom_dos

deepmd/tf/model/ener.py

@@ -193,6 +193,9 @@ def build(
             t_mt = tf.constant(self.model_type, name="model_type", dtype=tf.string)
             t_ver = tf.constant(MODEL_VERSION, name="model_version", dtype=tf.string)
 
+            # Initialize out_bias and out_std for energy models
+            self.init_out_stat(suffix=suffix)
+
             if self.srtab is not None:
                 tab_info, tab_data = self.srtab.get()
                 self.tab_info = tf.get_variable(
@@ -253,6 +256,10 @@ def build(
         atom_ener = self.fitting.build(
             dout, natoms, input_dict, reuse=reuse, suffix=suffix
         )
+
+        # Apply out_bias and out_std directly to atom energy
+        atom_ener = self._apply_out_bias_std(atom_ener, atype, natoms, coord)
+
         self.atom_ener = atom_ener
 
         if self.srtab is not None:

deepmd/tf/model/model.py

@@ -708,6 +708,63 @@ def __init__(
         else:
             self.typeebd = None
 
+        # Initialize out_bias and out_std storage
+        self.out_bias = None
+        self.out_std = None
+
+    def init_variables(
+        self,
+        graph: tf.Graph,
+        graph_def: tf.GraphDef,
+        model_type: str = "original_model",
+        suffix: str = "",
+    ) -> None:
+        """Init the model variables with the given frozen model.
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        model_type : str
+            the type of the model
+        suffix : str
+            suffix to name scope
+        """
+        from deepmd.tf.utils.errors import (
+            GraphWithoutTensorError,
+        )
+        from deepmd.tf.utils.graph import (
+            get_tensor_by_name_from_graph,
+        )
+
+        # Initialize descriptor and fitting variables
+        self.descrpt.init_variables(graph, graph_def, suffix=suffix)
+        self.fitting.init_variables(graph, graph_def, suffix=suffix)
+        if (
+            self.typeebd is not None
+            and self.typeebd.type_embedding_net_variables is None
+        ):
+            self.typeebd.init_variables(graph, graph_def, suffix=suffix)
+
+        # Try to load out_bias and out_std from the graph
+        try:
+            self.out_bias = get_tensor_by_name_from_graph(
+                graph, f"model_attr{suffix}/t_out_bias"
+            )
+        except GraphWithoutTensorError:
+            # For compatibility, create default out_bias if not found
+            pass
+
+        try:
+            self.out_std = get_tensor_by_name_from_graph(
+                graph, f"model_attr{suffix}/t_out_std"
+            )
+        except GraphWithoutTensorError:
+            # For compatibility, create default out_std if not found
+            pass
+
     def enable_mixed_precision(self, mixed_prec: dict) -> None:
         """Enable mixed precision for the model.
 
@@ -762,6 +819,114 @@ def get_ntypes(self) -> int:
         """Get the number of types."""
         return self.ntypes
 
+    def init_out_stat(self, suffix: str = "") -> None:
+        """Initialize the output bias and std variables."""
+        ntypes = self.get_ntypes()
+
+        # Determine output dimension based on model type instead of fitting type
+        if hasattr(self, "model_type"):
+            model_type = self.model_type
+        else:
+            # Fallback to fitting type for compatibility
+            model_type = getattr(self.fitting, "model_type", "ener")
+
+        if model_type == "ener":
+            dim_out = 1
+        elif model_type in ["dipole", "polar"]:
+            dim_out = 3
+        elif model_type == "dos":
+            dim_out = getattr(self.fitting, "numb_dos", 1)
+        else:
+            dim_out = 1
+
+        # Initialize out_bias and out_std as numpy arrays, preserving existing values if set
+        if hasattr(self, "out_bias") and self.out_bias is not None:
+            out_bias_data = self.out_bias.copy()
+        else:
+            out_bias_data = np.zeros(
+                [1, ntypes, dim_out], dtype=GLOBAL_NP_FLOAT_PRECISION
+            )
+
+        if hasattr(self, "out_std") and self.out_std is not None:
+            out_std_data = self.out_std.copy()
+        else:
+            out_std_data = np.ones(
+                [1, ntypes, dim_out], dtype=GLOBAL_NP_FLOAT_PRECISION
+            )
+
+        # Create TensorFlow variables
+        with tf.variable_scope("model_attr" + suffix, reuse=tf.AUTO_REUSE):
+            self.t_out_bias = tf.get_variable(
+                "t_out_bias",
+                out_bias_data.shape,
+                dtype=GLOBAL_TF_FLOAT_PRECISION,
+                trainable=False,
+                initializer=tf.constant_initializer(out_bias_data),
+            )
+            self.t_out_std = tf.get_variable(
+                "t_out_std",
+                out_std_data.shape,
+                dtype=GLOBAL_TF_FLOAT_PRECISION,
+                trainable=False,
+                initializer=tf.constant_initializer(out_std_data),
+            )
+
+        # Store as instance variables for access
+        self.out_bias = out_bias_data
+        self.out_std = out_std_data
+
+    def _apply_out_bias_std(self, output, atype, natoms, coord, selected_atype=None):
+        """Apply output bias and standard deviation to the model output.
+
+        Parameters
+        ----------
+        output : tf.Tensor
+            The model output tensor
+        atype : tf.Tensor
+            Atom types with shape [nframes, nloc]
+        natoms : list[int]
+            Number of atoms [nloc, ntypes, ...]
+        coord : tf.Tensor
+            Coordinates for getting nframes
+        selected_atype : tf.Tensor, optional
+            Selected atom types for tensor models. If None, uses all atoms.
+
+        Returns
+        -------
+        tf.Tensor
+            Output with bias and std applied
+        """
+        nframes = tf.shape(coord)[0]
+
+        if selected_atype is not None:
+            # For tensor models (dipole, polar) with selected atoms
+            natomsel = tf.shape(selected_atype)[1]
+            nout = self.get_out_size()  # Use the model's output size method
+            output_reshaped = tf.reshape(output, [nframes, natomsel, nout])
+            atype_for_gather = selected_atype
+        else:
+            # For energy and DOS models with all atoms
+            nloc = natoms[0]
+            if hasattr(self, "numb_dos"):
+                # DOS model: output shape [nframes * nloc * numb_dos]
+                nout = self.numb_dos
+                output_reshaped = tf.reshape(output, [nframes, nloc, nout])
+            else:
+                # Energy model: output shape [nframes * nloc]
+                nout = 1
+                output_reshaped = tf.reshape(output, [nframes, nloc, 1])
+            atype_for_gather = tf.reshape(atype, [nframes, nloc])
+
+        # Get bias and std for each atom type
+        bias_per_atom = tf.gather(self.t_out_bias[0], atype_for_gather)
+        std_per_atom = tf.gather(self.t_out_std[0], atype_for_gather)
+
+        # Apply bias and std: output = output * std + bias
+        output_reshaped = output_reshaped * std_per_atom + bias_per_atom
+
+        # Reshape back to original shape
+        return tf.reshape(output_reshaped, tf.shape(output))
+
     @classmethod
     def update_sel(
         cls,
@@ -820,25 +985,7 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor":
         data = data.copy()
         check_version_compatibility(data.pop("@version", 2), 2, 1)
         descriptor = Descriptor.deserialize(data.pop("descriptor"), suffix=suffix)
-        if data["fitting"].get("@variables", {}).get("bias_atom_e") is not None:
-            # careful: copy each level and don't modify the input array,
-            # otherwise it will affect the original data
-            # deepcopy is not used for performance reasons
-            data["fitting"] = data["fitting"].copy()
-            data["fitting"]["@variables"] = data["fitting"]["@variables"].copy()
-            if (
-                int(np.any(data["fitting"]["@variables"]["bias_atom_e"]))
-                + int(np.any(data["@variables"]["out_bias"]))
-                > 1
-            ):
-                raise ValueError(
-                    "fitting/@variables/bias_atom_e and @variables/out_bias should not be both non-zero"
-                )
-            data["fitting"]["@variables"]["bias_atom_e"] = data["fitting"][
-                "@variables"
-            ]["bias_atom_e"] + data["@variables"]["out_bias"].reshape(
-                data["fitting"]["@variables"]["bias_atom_e"].shape
-            )
+        # bias_atom_e and out_bias are now completely independent - no conversion needed
         fitting = Fitting.deserialize(data.pop("fitting"), suffix=suffix)
         # pass descriptor type embedding to model
         if descriptor.explicit_ntypes:
@@ -853,14 +1000,23 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor":
             raise NotImplementedError("pair_exclude_types is not supported")
         data.pop("rcond", None)
         data.pop("preset_out_bias", None)
-        data.pop("@variables", None)
+        # Extract out_bias and out_std from variables before removing them
+        variables = data.pop("@variables", {})
+        out_bias = variables.get("out_bias", None)
+        out_std = variables.get("out_std", None)
         # END    not supported keys
-        return cls(
+        model = cls(
             descriptor=descriptor,
             fitting_net=fitting,
             type_embedding=type_embedding,
             **data,
         )
+        # Restore out_bias and out_std if they exist
+        if out_bias is not None:
+            model.out_bias = out_bias
+        if out_std is not None:
+            model.out_std = out_std
+        return model
 
     def serialize(self, suffix: str = "") -> dict:
         """Serialize the model.
@@ -886,8 +1042,41 @@ def serialize(self, suffix: str = "") -> dict:
             raise NotImplementedError("spin is not supported")
 
         ntypes = len(self.get_type_map())
-        dict_fit = self.fitting.serialize(suffix=suffix)
-        if dict_fit.get("@variables", {}).get("bias_atom_e") is not None:
+
+        # Try to serialize fitting, with fallback for uninitialized variables
+        try:
+            dict_fit = self.fitting.serialize(suffix=suffix)
+        except (AttributeError, TypeError):
+            # Fallback: create a minimal dict_fit with just dim_out
+            from deepmd.tf.fit.dipole import (
+                DipoleFittingSeA,
+            )
+            from deepmd.tf.fit.dos import (
+                DOSFitting,
+            )
+            from deepmd.tf.fit.ener import (
+                EnerFitting,
+            )
+            from deepmd.tf.fit.polar import (
+                PolarFittingSeA,
+            )
+
+            if isinstance(self.fitting, EnerFitting):
+                dim_out = 1
+            elif isinstance(self.fitting, (DipoleFittingSeA, PolarFittingSeA)):
+                dim_out = 3
+            elif isinstance(self.fitting, DOSFitting):
+                dim_out = getattr(self.fitting, "numb_dos", 1)
+            else:
+                dim_out = 1
+
+            dict_fit = {"dim_out": dim_out, "@variables": {}}
+
+        # Use the actual out_bias and out_std if they exist, otherwise create defaults
+        if self.out_bias is not None:
+            out_bias = self.out_bias.copy()
+        elif dict_fit.get("@variables", {}).get("bias_atom_e") is not None:
+            # Fallback to converting bias_atom_e for backward compatibility
             out_bias = dict_fit["@variables"]["bias_atom_e"].reshape(
                 [1, ntypes, dict_fit["dim_out"]]
             )
@@ -898,6 +1087,13 @@ def serialize(self, suffix: str = "") -> dict:
             out_bias = np.zeros(
                 [1, ntypes, dict_fit["dim_out"]], dtype=GLOBAL_NP_FLOAT_PRECISION
             )
+
+        if self.out_std is not None:
+            out_std = self.out_std.copy()
+        else:
+            out_std = np.ones(
+                [1, ntypes, dict_fit["dim_out"]], dtype=GLOBAL_NP_FLOAT_PRECISION
+            )
         return {
             "@class": "Model",
             "type": "standard",
@@ -912,7 +1108,7 @@ def serialize(self, suffix: str = "") -> dict:
             "preset_out_bias": None,
             "@variables": {
                 "out_bias": out_bias,
-                "out_std": np.ones([1, ntypes, dict_fit["dim_out"]]),  # pylint: disable=no-explicit-dtype
+                "out_std": out_std,
             },
         }
 

deepmd/tf/model/tensor.py

@@ -126,6 +126,9 @@ def build(
             t_ver = tf.constant(MODEL_VERSION, name="model_version", dtype=tf.string)
             t_od = tf.constant(self.get_out_size(), name="output_dim", dtype=tf.int32)
 
+            # Initialize out_bias and out_std for tensor models (dipole/polar)
+            self.init_out_stat(suffix=suffix)
+
         natomsel = sum(natoms[2 + type_i] for type_i in self.get_sel_type())
         nout = self.get_out_size()
 
@@ -164,6 +167,12 @@ def build(
         output = self.fitting.build(
             dout, rot_mat, natoms, input_dict, reuse=reuse, suffix=suffix
         )
+
+        # Apply out_bias and out_std directly to tensor output
+        atype_selected = self._get_selected_atype(atype, natoms)
+        output = self._apply_out_bias_std(
+            output, atype, natoms, coord, selected_atype=atype_selected
+        )
         framesize = nout if "global" in self.model_type else natomsel * nout
         output = tf.reshape(
             output, [-1, framesize], name="o_" + self.model_type + suffix
@@ -206,6 +215,24 @@ def build(
 
         return model_dict
 
+    def _get_selected_atype(self, atype, natoms):
+        """Get atom types for selected atoms only (matching tensor model selection)."""
+        # For tensor models, the fitting output corresponds to selected atom types
+        # atype shape: [nframes, nloc]
+        # We need to extract atom types that match the natomsel count
+
+        # Simplified approach: take the first natomsel atoms from each frame
+        # This works because natoms and descriptor arrangement should be consistent
+        nframes = tf.shape(atype)[0]
+        selected_types = self.get_sel_type()
+        natomsel = sum(natoms[2 + type_i] for type_i in selected_types)
+
+        # Take the first natomsel atoms from each frame
+        # This assumes the atom ordering is consistent with how fitting produces output
+        atype_selected = atype[:, :natomsel]  # [nframes, natomsel]
+
+        return atype_selected
+
     def init_variables(
         self,
         graph: tf.Graph,