feat(pt): Implement type embedding compression for two-side mode for se_atten

- Added functionality to enable type embedding compression in the  class, specifically for two-side mode.
- Introduced a new method  to precompute type embedding network outputs for all type pair combinations, optimizing inference by using precomputed values.
- Updated the  method to utilize precomputed embeddings when compression is enabled, improving performance during inference.

This enhancement allows for more efficient handling of type embeddings in the descriptor model.

Author: OutisLi

deepmd/pt/model/descriptor/dpa1.py

@@ -645,6 +645,12 @@ def enable_compression(
         self.se_atten.enable_compression(
             self.table.data, self.table_config, self.lower, self.upper
         )
+
+        # Enable type embedding compression only for two-side mode
+        # TODO: why not enable for one-side mode? (do not consider this for now)
+        if not self.se_atten.type_one_side:
+            self.se_atten.enable_type_embedding_compression(self.type_embedding)
+
         self.compress = True
 
     def forward(

deepmd/pt/model/descriptor/se_atten.py

@@ -27,6 +27,9 @@
     MLPLayer,
     NetworkCollection,
 )
+from deepmd.pt.model.network.network import (
+    TypeEmbedNet,
+)
 from deepmd.pt.utils import (
     env,
 )
@@ -281,6 +284,9 @@ def __init__(
         self.compress_data = nn.ParameterList(
             [nn.Parameter(torch.zeros(0, dtype=self.prec, device=env.DEVICE))]
         )
+        # For type embedding compression (strip mode, two-side only)
+        self.compress_type_embd = False
+        self.two_side_embd_data = None
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
@@ -447,6 +453,50 @@ def enable_compression(
         self.compress_data[0] = table_data[net].to(device=env.DEVICE, dtype=self.prec)
         self.compress = True
 
+    def enable_type_embedding_compression(
+        self, type_embedding_net: TypeEmbedNet
+    ) -> None:
+        """Enable type embedding compression for strip mode (two-side only).
+
+        This method precomputes the type embedding network outputs for all possible
+        type pairs, following the same approach as TF backend's compression:
+
+        TF approach:
+        1. get_two_side_type_embedding(): creates (ntypes+1)^2 type pair combinations
+        2. make_data(): applies embedding network to get precomputed outputs
+        3. In forward: lookup precomputed values instead of real-time computation
+
+        PyTorch implementation:
+        - Precomputes all (ntypes+1)^2 type pair embedding network outputs
+        - Stores in buffer for proper serialization and device management
+        - Uses lookup during inference to avoid redundant computations
+
+        Parameters
+        ----------
+        type_embedding_net : TypeEmbedNet
+            The type embedding network that provides get_full_embedding() method
+        """
+        with torch.no_grad():
+            # Get full type embedding: (ntypes+1) x t_dim
+            full_embd = type_embedding_net.get_full_embedding(env.DEVICE)
+            nt, t_dim = full_embd.shape
+
+            # Create all type pair combinations [neighbor, center]
+            # for a fixed row i, all columns j have different neighbor types
+            embd_nei = full_embd.view(1, nt, t_dim).expand(nt, nt, t_dim)
+            # for a fixed row i, all columns j share the same center type i
+            embd_center = full_embd.view(nt, 1, t_dim).expand(nt, nt, t_dim)
+            two_side_embd = torch.cat([embd_nei, embd_center], dim=-1).reshape(
+                -1, t_dim * 2
+            )
+
+            # Apply strip embedding network and store
+            # index logic: index = center_type * nt + neighbor_type
+            self.two_side_embd_data = self.filter_layers_strip.networks[0](
+                two_side_embd
+            ).detach()
+            self.compress_type_embd = True
+
     def forward(
         self,
         nlist: torch.Tensor,
@@ -605,7 +655,12 @@ def forward(
                 two_side_type_embedding = torch.cat(
                     [type_embedding_nei, type_embedding_center], -1
                 ).reshape(-1, nt * 2)
-                tt_full = self.filter_layers_strip.networks[0](two_side_type_embedding)
+                if self.compress_type_embd and self.two_side_embd_data is not None:
+                    tt_full = self.two_side_embd_data
+                else:
+                    tt_full = self.filter_layers_strip.networks[0](
+                        two_side_type_embedding
+                    )
                 # (nf x nl x nnei) x ng
                 gg_t = torch.gather(tt_full, dim=0, index=idx)
             # (nf x nl) x nnei x ng