refactor(tf): consolidate out_bias/out_std application into shared method

Co-authored-by: njzjz <[email protected]>

Author: Copilot

deepmd/tf/model/dos.py

@@ -186,28 +186,7 @@ def build(
         )
 
         # Apply out_bias and out_std directly to DOS output
-        # atom_dos shape: [nframes * nloc * numb_dos] for DOS models
-        # t_out_bias shape: [1, ntypes, numb_dos], t_out_std shape: [1, ntypes, numb_dos]
-        if hasattr(self, "t_out_bias") and hasattr(self, "t_out_std"):
-            nframes = tf.shape(coord)[0]
-            nloc = natoms[0]
-            # Reshape atom_dos to [nframes, nloc, numb_dos] for bias/std application
-            atom_dos_reshaped = tf.reshape(atom_dos, [nframes, nloc, self.numb_dos])
-
-            # Get bias and std for each atom type: [nframes, nloc, numb_dos]
-            atype_flat = tf.reshape(atype, [nframes, nloc])
-            bias_per_atom = tf.gather(
-                self.t_out_bias[0], atype_flat
-            )  # [nframes, nloc, numb_dos]
-            std_per_atom = tf.gather(
-                self.t_out_std[0], atype_flat
-            )  # [nframes, nloc, numb_dos]
-
-            # Apply bias and std: dos = dos * std + bias
-            atom_dos_reshaped = atom_dos_reshaped * std_per_atom + bias_per_atom
-
-            # Reshape back to original shape
-            atom_dos = tf.reshape(atom_dos_reshaped, tf.shape(atom_dos))
+        atom_dos = self._apply_out_bias_std(atom_dos, atype, natoms, coord)
 
         self.atom_dos = atom_dos
 

deepmd/tf/model/ener.py

@@ -258,29 +258,7 @@ def build(
         )
 
         # Apply out_bias and out_std directly to atom energy
-        # atom_ener shape: [nframes * nloc] (for energy models, dim_out=1)
-        # t_out_bias shape: [1, ntypes, 1], t_out_std shape: [1, ntypes, 1]
-        # atype shape: [nframes, nloc]
-        if hasattr(self, "t_out_bias") and hasattr(self, "t_out_std"):
-            # Reshape atom_ener to [nframes, nloc, 1] to match bias/std application
-            nframes = tf.shape(coord)[0]
-            nloc = natoms[0]
-            atom_ener_reshaped = tf.reshape(atom_ener, [nframes, nloc, 1])
-
-            # Get bias and std for each atom type: [nframes, nloc, 1]
-            atype_flat = tf.reshape(atype, [nframes, nloc])
-            bias_per_atom = tf.gather(
-                self.t_out_bias[0], atype_flat
-            )  # [nframes, nloc, 1]
-            std_per_atom = tf.gather(
-                self.t_out_std[0], atype_flat
-            )  # [nframes, nloc, 1]
-
-            # Apply bias and std: energy = energy * std + bias
-            atom_ener_reshaped = atom_ener_reshaped * std_per_atom + bias_per_atom
-
-            # Reshape back to original shape
-            atom_ener = tf.reshape(atom_ener_reshaped, tf.shape(atom_ener))
+        atom_ener = self._apply_out_bias_std(atom_ener, atype, natoms, coord)
 
         self.atom_ener = atom_ener
 

deepmd/tf/model/model.py

@@ -823,33 +823,21 @@ def init_out_stat(self, suffix: str = "") -> None:
         """Initialize the output bias and std variables."""
         ntypes = self.get_ntypes()
 
-        # Get output dimension from fitting serialization, with fallback
-        try:
-            dict_fit = self.fitting.serialize(suffix=suffix)
-            dim_out = dict_fit.get("dim_out", 1)
-        except (AttributeError, TypeError):
-            # Fallback to default dimensions for different fitting types
-            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
+        # 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:
@@ -887,31 +875,57 @@ def init_out_stat(self, suffix: str = "") -> None:
         self.out_bias = out_bias_data
         self.out_std = out_std_data
 
-    def get_out_bias(self) -> np.ndarray:
-        """Get the output bias."""
-        return self.out_bias
-
-    def get_out_std(self) -> np.ndarray:
-        """Get the output standard deviation."""
-        return self.out_std
-
-    def set_out_bias(self, out_bias: np.ndarray) -> None:
-        """Set the output bias."""
-        self.out_bias = out_bias
-        if hasattr(self, "t_out_bias"):
-            # Note: TensorFlow variable assignment would require a session context in TF 1.x
-            # For TF 2.x, the variable assignment happens differently
-            # Here we just update the numpy array, and TF variables are updated when rebuilt
-            pass
+    def _apply_out_bias_std(self, output, atype, natoms, coord, selected_atype=None):
+        """Apply output bias and standard deviation to the model output.
 
-    def set_out_std(self, out_std: np.ndarray) -> None:
-        """Set the output standard deviation."""
-        self.out_std = out_std
-        if hasattr(self, "t_out_std"):
-            # Note: TensorFlow variable assignment would require a session context in TF 1.x
-            # For TF 2.x, the variable assignment happens differently
-            # Here we just update the numpy array, and TF variables are updated when rebuilt
-            pass
+        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(

deepmd/tf/model/tensor.py

@@ -169,29 +169,10 @@ def build(
         )
 
         # Apply out_bias and out_std directly to tensor output
-        # output shape: [nframes * natomsel * nout] for tensor models
-        # t_out_bias shape: [1, ntypes, nout], t_out_std shape: [1, ntypes, nout]
-        if hasattr(self, "t_out_bias") and hasattr(self, "t_out_std"):
-            nframes = tf.shape(coord)[0]
-            # Reshape output to [nframes, natomsel, nout] for bias/std application
-            output_reshaped = tf.reshape(output, [nframes, natomsel, nout])
-
-            # Get atom types for selected atoms only (matching natomsel)
-            atype_selected = self._get_selected_atype(atype, natoms)
-
-            # Get bias and std for each selected atom type: [nframes, natomsel, nout]
-            bias_per_atom = tf.gather(
-                self.t_out_bias[0], atype_selected
-            )  # [nframes, natomsel, nout]
-            std_per_atom = tf.gather(
-                self.t_out_std[0], atype_selected
-            )  # [nframes, natomsel, nout]
-
-            # Apply bias and std: output = output * std + bias
-            output_reshaped = output_reshaped * std_per_atom + bias_per_atom
-
-            # Reshape back to original shape
-            output = tf.reshape(output_reshaped, tf.shape(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