feat(dpa1): add compressible l=2 moments across PT backends

deepmd/dpmodel/descriptor/se_atten_v2.py

@@ -122,6 +122,9 @@ class DescrptSeAttenV2(DescrptDPA1):
             A list of strings. Give the name to each type of atoms.
     seed : int, Optional
             Random seed for initializing the network parameters.
+    lmax : int
+            Maximum angular degree of the Cartesian moment basis. Supported
+            values are 1 through 4.
     """
 
     def __init__(
@@ -158,6 +161,7 @@ def __init__(
         type_map: list[str] | None = None,
         # consistent with argcheck, not used though
         seed: int | list[int] | None = None,
+        lmax: int = 1,
     ) -> None:
         DescrptDPA1.__init__(
             self,
@@ -195,6 +199,7 @@ def __init__(
             type_map=type_map,
             # consistent with argcheck, not used though
             seed=seed,
+            lmax=lmax,
         )
         self.compress = False
 
@@ -204,7 +209,7 @@ def serialize(self) -> dict:
         data = {
             "@class": "Descriptor",
             "type": "se_atten_v2",
-            "@version": 3 if self.compress else 2,
+            "@version": 4 if obj.lmax != 1 else (3 if self.compress else 2),
             "rcut": obj.rcut,
             "rcut_smth": obj.rcut_smth,
             "sel": obj.sel,
@@ -246,6 +251,11 @@ def serialize(self) -> dict:
             "trainable": self.trainable,
             "spin": None,
         }
+        if obj.lmax != 1:
+            data["lmax"] = obj.lmax
+            data["@variables"]["degree_gain_raw"] = to_numpy_array(
+                obj.adam_degree_gain_raw
+            )
         if self.compress:
             type_embd_data = (
                 self.type_embd_data
@@ -282,7 +292,7 @@ def serialize(self) -> dict:
     def deserialize(cls, data: dict) -> "DescrptSeAttenV2":
         """Deserialize from dict."""
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 3, 1)
+        check_version_compatibility(data.pop("@version"), 4, 1)
         data.pop("@class")
         data.pop("type")
         variables = data.pop("@variables")
@@ -295,10 +305,16 @@ def deserialize(cls, data: dict) -> "DescrptSeAttenV2":
         # compat with version 1
         if "use_tebd_bias" not in data:
             data["use_tebd_bias"] = True
+        data.setdefault("lmax", 1)
         obj = cls(**data)
 
         obj.se_atten["davg"] = variables["davg"]
         obj.se_atten["dstd"] = variables["dstd"]
+        if obj.se_atten.lmax > 1:
+            obj.se_atten.adam_degree_gain_raw = np.asarray(
+                variables["degree_gain_raw"],
+                dtype=PRECISION_DICT[obj.se_atten.precision],
+            )
         obj.se_atten.embeddings = NetworkCollection.deserialize(embeddings)
         obj.se_atten.embeddings_strip = NetworkCollection.deserialize(embeddings_strip)
         obj.type_embedding = TypeEmbedNet.deserialize(type_embedding)

deepmd/kernels/cuda/dpa1/canonical.py

@@ -70,6 +70,7 @@ def _forward_fake(
     type_embedding: torch.Tensor,
     average: torch.Tensor,
     inverse_stddev: torch.Tensor,
+    degree_gain: torch.Tensor,
     table: torch.Tensor,
     gate_table: torch.Tensor,
     type_one_side: int,
@@ -86,12 +87,14 @@ def _forward_fake(
     rcut_smooth: float,
     protection: float,
     neighbors: float,
+    basis_dim: int,
 ) -> tuple[torch.Tensor, ...]:
     del (
         source,
         destination_row_ptr,
         average,
         inverse_stddev,
+        degree_gain,
         gate_table,
         type_one_side,
         smooth,
@@ -117,7 +120,7 @@ def _forward_fake(
             width,
             3,
         ),
-        edge_vec.new_empty(node_count, 4, width),
+        edge_vec.new_empty(node_count, basis_dim, width),
     )
 
 
@@ -131,6 +134,7 @@ def _backward_fake(
     atype: torch.Tensor,
     average: torch.Tensor,
     inverse_stddev: torch.Tensor,
+    degree_gain: torch.Tensor,
     table: torch.Tensor,
     gate_table: torch.Tensor,
     type_one_side: int,
@@ -155,6 +159,7 @@ def _backward_fake(
         atype,
         average,
         inverse_stddev,
+        degree_gain,
         table,
         gate_table,
         type_one_side,
@@ -218,9 +223,9 @@ def _cpu_forward(*args: Any) -> tuple[torch.Tensor, ...]:
         edge_mask,
         destination_order,
         destination_row_ptr,
-        *tail[:11],
+        *tail[:12],
         True,
-        *tail[11:],
+        *tail[12:],
     )
 
 
@@ -245,9 +250,9 @@ def _cpu_backward(*args: Any) -> torch.Tensor:
         edge_mask,
         destination_order,
         destination_row_ptr,
-        *tail[:8],
+        *tail[:9],
         True,
-        *tail[8:],
+        *tail[9:],
     )
 
 
@@ -329,6 +334,11 @@ def dpa1_canonical_compress_energy_force(
     compress_data = desc.compress_data[0].contiguous()
     gate_table = desc.type_embd_data.contiguous()
     inverse_stddev = torch.reciprocal(se.stddev[:, 0, :]).contiguous()
+    degree_gain = (
+        se.adam_degree_gain_raw.to(torch.float32).contiguous()
+        if se.adam_degree_gain_raw is not None
+        else compress_data.new_empty(0)
+    )
     from torch.fx.experimental.proxy_tensor import (
         disable_proxy_modes_tracing,
     )
@@ -346,6 +356,7 @@ def dpa1_canonical_compress_energy_force(
         type_embedding,
         se.mean[:, 0, :].contiguous(),
         inverse_stddev,
+        degree_gain,
         compress_data,
         gate_table,
         int(se.type_one_side),
@@ -362,6 +373,7 @@ def dpa1_canonical_compress_energy_force(
         float(se.rcut_smth),
         float(se.env_protection),
         float(se.nnei),
+        (int(se.lmax) + 1) ** 2,
     )
 
     *hidden, head = fit.nets[0].layers
@@ -426,6 +438,7 @@ def dpa1_canonical_compress_energy_force(
         atype,
         se.mean[:, 0, :].contiguous(),
         inverse_stddev,
+        degree_gain,
         compress_data,
         gate_table,
         int(se.type_one_side),