feat(pt_expt): add dp compress support for pt_expt backend (#5323)

## Summary

- Add model compression (embedding net tabulation) for the pt_expt
backend, matching the existing pt backend capability
- Compressed models replace embedding net forward passes with polynomial
lookup tables via C++ custom ops (`tabulate_fusion_se_*`), significantly
speeding up inference
- Support all compressible descriptors: `se_e2_a`, `se_r`, `se_t`,
`se_t_tebd`, `dpa1`, `se_atten_v2`, `dpa2` (hybrid delegates
automatically)

### Key changes

**Infrastructure:**
- `deepmd/pt_expt/utils/tabulate_ops.py` — Register
`torch.library.register_fake` for all 5 custom ops to enable
`torch.export`/`make_fx` tracing through compressed forward paths
- `deepmd/pt_expt/utils/tabulate.py` — `DPTabulate` subclass that
detects descriptor type via serialized data (avoids `isinstance` checks
against pt-specific classes)
- `deepmd/pt_expt/entrypoints/compress.py` — Entry point: load `.pte` →
deserialize → `enable_compression()` → re-export `.pte`

**Descriptors:** Each gets `enable_compression()` + `@cast_precision`
`call()` override with compressed branch using the appropriate custom
op.

**dpmodel — compression state serialization (breaking version bumps):**

The pt_expt backend persists models via `serialize()` → `model.json` →
`deserialize()` (the `.pte` format), unlike pt/tf which use native
framework save mechanisms (torch.jit.save / tf.saved_model) that capture
the full runtime state. This means compression state (tabulated
polynomial coefficients, precomputed type embeddings) must survive the
serialize/deserialize round-trip for compressed `.pte` models to work.

Each compressible descriptor's serialization version is bumped when the
model is compressed. **Uncompressed models continue to use the old
version**, so there is no breakage for existing uncompressed model
files. All backends (pt, pd, tf) accept the new version in
`deserialize()` and simply ignore the `"compress"` key.

| Descriptor | Version bump | Added fields |
|---|---|---|
| `se_e2_a` | 2 → 3 | `compress_data`, `compress_info` |
| `se_r` | 2 → 3 | `compress_data`, `compress_info` |
| `se_t` | 2 → 3 | `compress_data`, `compress_info` |
| `se_t_tebd` | 1 → 2 | `compress_data`, `compress_info`,
`type_embd_data` |
| `dpa1` | 2 → 3 | `type_embd_data`, `geo_compress`,
`compress_data`/`info` (if geo) |
| `se_atten_v2` | 2 → 3 | `type_embd_data`, `geo_compress`,
`compress_data`/`info` (if geo) |
| `dpa2` | 3 → 4 | compress dict inside `repinit_variable` |

**dpmodel:** Initialize `self.compress = False` in all descriptor
`__init__` methods.

## Test plan

- [x] `source/tests/pt_expt/model/test_model_compression.py` —
end-to-end compress → serialize → deserialize → eval
- [x] `source/tests/pt_expt/descriptor/` — compressed forward,
consistency, exportable, make_fx tests for all descriptors
- [x] `source/tests/consistent/descriptor/` — cross-backend consistency
tests pass with bumped versions

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added descriptor compression functionality to reduce model size and
optimize memory usage during inference.
* Introduced `compress` CLI command to enable tabulated embedding
optimization on frozen trained models.
* Enhanced descriptor serialization with improved version compatibility
across multiple backends.

* **Tests**
* Added comprehensive test coverage for compressed descriptor forward
passes and model compression workflows.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

deepmd/dpmodel/descriptor/dpa1.py

@@ -344,6 +344,7 @@ def __init__(
         self.concat_output_tebd = concat_output_tebd
         self.trainable = trainable
         self.precision = precision
+        self.compress = False
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
@@ -557,7 +558,7 @@ def serialize(self) -> dict:
         data = {
             "@class": "Descriptor",
             "type": "dpa1",
-            "@version": 2,
+            "@version": 3 if self.compress else 2,
             "rcut": obj.rcut,
             "rcut_smth": obj.rcut_smth,
             "sel": obj.sel,
@@ -602,20 +603,36 @@ def serialize(self) -> dict:
         }
         if obj.tebd_input_mode in ["strip"]:
             data.update({"embeddings_strip": obj.embeddings_strip.serialize()})
+        if self.compress:
+            compress_dict: dict = {
+                "@variables": {
+                    "type_embd_data": to_numpy_array(self.type_embd_data),
+                },
+                "geo_compress": self.geo_compress,
+            }
+            if self.geo_compress:
+                compress_dict["@variables"]["compress_data"] = [
+                    to_numpy_array(d) for d in self.compress_data
+                ]
+                compress_dict["@variables"]["compress_info"] = [
+                    to_numpy_array(i) for i in self.compress_info
+                ]
+            data["compress"] = compress_dict
         return data
 
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptDPA1":
         """Deserialize from dict."""
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 2, 1)
+        check_version_compatibility(data.pop("@version"), 3, 1)
         data.pop("@class")
         data.pop("type")
         variables = data.pop("@variables")
         embeddings = data.pop("embeddings")
         type_embedding = data.pop("type_embedding")
         attention_layers = data.pop("attention_layers")
         env_mat = data.pop("env_mat")
+        compress = data.pop("compress", None)
         tebd_input_mode = data["tebd_input_mode"]
         if tebd_input_mode in ["strip"]:
             embeddings_strip = data.pop("embeddings_strip")
@@ -637,8 +654,20 @@ def deserialize(cls, data: dict) -> "DescrptDPA1":
         obj.se_atten.dpa1_attention = NeighborGatedAttention.deserialize(
             attention_layers
         )
+        if compress is not None:
+            obj._load_compress_data(compress)
         return obj
 
+    def _load_compress_data(self, compress: dict) -> None:
+        """Load compression state from serialized data."""
+        variables = compress["@variables"]
+        self.type_embd_data = variables["type_embd_data"]
+        self.geo_compress = compress.get("geo_compress", False)
+        if self.geo_compress:
+            self.compress_data = variables["compress_data"]
+            self.compress_info = variables["compress_info"]
+        self.compress = True
+
     @classmethod
     def update_sel(
         cls,

deepmd/dpmodel/descriptor/dpa2.py

@@ -596,6 +596,7 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
         self.rcut_smth = self.repinit.get_rcut_smth()
         self.trainable = trainable
         self.add_tebd_to_repinit_out = add_tebd_to_repinit_out
+        self.compress = False
 
         self.repinit_out_dim = self.repinit.dim_out
         if self.repinit_args.use_three_body:
@@ -938,7 +939,7 @@ def serialize(self) -> dict:
         data = {
             "@class": "Descriptor",
             "type": "dpa2",
-            "@version": 3,
+            "@version": 4 if self.compress else 3,
             "ntypes": self.ntypes,
             "repinit_args": self.repinit_args.serialize(),
             "repformer_args": self.repformer_args.serialize(),
@@ -973,6 +974,21 @@ def serialize(self) -> dict:
             repinit_variable.update(
                 {"embeddings_strip": repinit.embeddings_strip.serialize()}
             )
+        if self.compress:
+            compress_dict: dict = {
+                "@variables": {
+                    "type_embd_data": to_numpy_array(self.type_embd_data),
+                },
+                "geo_compress": self.geo_compress,
+            }
+            if self.geo_compress:
+                compress_dict["@variables"]["compress_data"] = [
+                    to_numpy_array(d) for d in self.compress_data
+                ]
+                compress_dict["@variables"]["compress_info"] = [
+                    to_numpy_array(i) for i in self.compress_info
+                ]
+            repinit_variable["compress"] = compress_dict
         repformers_variable = {
             "g2_embd": repformers.g2_embd.serialize(),
             "repformer_layers": [layer.serialize() for layer in repformers.layers],
@@ -1016,7 +1032,7 @@ def serialize(self) -> dict:
     def deserialize(cls, data: dict) -> "DescrptDPA2":
         data = data.copy()
         version = data.pop("@version")
-        check_version_compatibility(version, 3, 1)
+        check_version_compatibility(version, 4, 1)
         data.pop("@class")
         data.pop("type")
         repinit_variable = data.pop("repinit_variable").copy()
@@ -1040,6 +1056,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA2":
         # compat with version 1
         if "use_tebd_bias" not in data:
             data["use_tebd_bias"] = True
+        compress = repinit_variable.pop("compress", None)
         obj = cls(**data)
         obj.type_embedding = TypeEmbedNet.deserialize(type_embedding)
         if add_tebd_to_repinit_out:
@@ -1089,8 +1106,20 @@ def deserialize(cls, data: dict) -> "DescrptDPA2":
         obj.repformers.layers = [
             RepformerLayer.deserialize(layer) for layer in repformer_layers
         ]
+        if compress is not None:
+            obj._load_compress_data(compress)
         return obj
 
+    def _load_compress_data(self, compress: dict) -> None:
+        """Load compression state from serialized data."""
+        variables = compress["@variables"]
+        self.type_embd_data = variables["type_embd_data"]
+        self.geo_compress = compress.get("geo_compress", False)
+        if self.geo_compress:
+            self.compress_data = variables["compress_data"]
+            self.compress_info = variables["compress_info"]
+        self.compress = True
+
     @classmethod
     def update_sel(
         cls,

deepmd/dpmodel/descriptor/se_atten_v2.py

@@ -196,14 +196,15 @@ def __init__(
             # consistent with argcheck, not used though
             seed=seed,
         )
+        self.compress = False
 
     def serialize(self) -> dict:
         """Serialize the descriptor to dict."""
         obj = self.se_atten
         data = {
             "@class": "Descriptor",
             "type": "se_atten_v2",
-            "@version": 2,
+            "@version": 3 if self.compress else 2,
             "rcut": obj.rcut,
             "rcut_smth": obj.rcut_smth,
             "sel": obj.sel,
@@ -245,13 +246,28 @@ def serialize(self) -> dict:
             "trainable": self.trainable,
             "spin": None,
         }
+        if self.compress:
+            compress_dict: dict = {
+                "@variables": {
+                    "type_embd_data": to_numpy_array(self.type_embd_data),
+                },
+                "geo_compress": self.geo_compress,
+            }
+            if self.geo_compress:
+                compress_dict["@variables"]["compress_data"] = [
+                    to_numpy_array(d) for d in self.compress_data
+                ]
+                compress_dict["@variables"]["compress_info"] = [
+                    to_numpy_array(i) for i in self.compress_info
+                ]
+            data["compress"] = compress_dict
         return data
 
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptSeAttenV2":
         """Deserialize from dict."""
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 2, 1)
+        check_version_compatibility(data.pop("@version"), 3, 1)
         data.pop("@class")
         data.pop("type")
         variables = data.pop("@variables")
@@ -260,6 +276,7 @@ def deserialize(cls, data: dict) -> "DescrptSeAttenV2":
         attention_layers = data.pop("attention_layers")
         data.pop("env_mat")
         embeddings_strip = data.pop("embeddings_strip")
+        compress = data.pop("compress", None)
         # compat with version 1
         if "use_tebd_bias" not in data:
             data["use_tebd_bias"] = True
@@ -273,4 +290,16 @@ def deserialize(cls, data: dict) -> "DescrptSeAttenV2":
         obj.se_atten.dpa1_attention = NeighborGatedAttention.deserialize(
             attention_layers
         )
+        if compress is not None:
+            obj._load_compress_data(compress)
         return obj
+
+    def _load_compress_data(self, compress: dict) -> None:
+        """Load compression state from serialized data."""
+        variables = compress["@variables"]
+        self.type_embd_data = variables["type_embd_data"]
+        self.geo_compress = compress.get("geo_compress", False)
+        if self.geo_compress:
+            self.compress_data = variables["compress_data"]
+            self.compress_info = variables["compress_info"]
+        self.compress = True

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -192,6 +192,7 @@ def __init__(
         self.precision = precision
         self.spin = spin
         self.type_map = type_map
+        self.compress = False
         # order matters, placed after the assignment of self.ntypes
         self.reinit_exclude(exclude_types)
 
@@ -514,10 +515,24 @@ def serialize(self) -> dict:
                 if embedding_idx in self.emask:
                     self.embeddings[embedding_idx].clear()
 
-        return {
+        # Serialization version history:
+        #   v2: original format
+        #   v3: added optional "compress" key for model compression state
+        #       (tabulated embedding-net polynomial coefficients).
+        #       Uncompressed models stay at v2 for full backward compatibility.
+        #
+        # Why serialize compression state here:
+        #   The pt and tf backends persist compression via their native
+        #   framework save mechanisms (torch.jit.save / tf.saved_model),
+        #   which capture the full runtime state including tabulated data.
+        #   The pt_expt backend uses serialize() -> model.json -> deserialize()
+        #   as the primary persistence path (.pte files), so compression state
+        #   must survive this round-trip.  All backends accept v3 in
+        #   deserialize() to allow cross-backend loading of compressed models.
+        data = {
             "@class": "Descriptor",
             "type": "se_e2_a",
-            "@version": 2,
+            "@version": 3 if self.compress else 2,
             "rcut": self.rcut,
             "rcut_smth": self.rcut_smth,
             "sel": self.sel,
@@ -541,24 +556,50 @@ def serialize(self) -> dict:
             },
             "type_map": self.type_map,
         }
+        if self.compress:
+            data["compress"] = {
+                "@variables": {
+                    "compress_data": [to_numpy_array(d) for d in self.compress_data],
+                    "compress_info": [to_numpy_array(i) for i in self.compress_info],
+                },
+            }
+        return data
 
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptSeA":
         """Deserialize from dict."""
         data = data.copy()
-        check_version_compatibility(data.pop("@version", 1), 2, 1)
+        check_version_compatibility(data.pop("@version", 1), 3, 1)
         data.pop("@class", None)
         data.pop("type", None)
         variables = data.pop("@variables")
         embeddings = data.pop("embeddings")
         env_mat = data.pop("env_mat")
+        compress = data.pop("compress", None)
         obj = cls(**data)
 
         obj["davg"] = variables["davg"]
         obj["dstd"] = variables["dstd"]
         obj.embeddings = NetworkCollection.deserialize(embeddings)
+        if compress is not None:
+            obj._load_compress_data(compress)
         return obj
 
+    def _load_compress_data(self, compress: dict) -> None:
+        """Load compression state from serialized data.
+
+        Parameters
+        ----------
+        compress : dict
+            Must contain "@variables" with "compress_data" (list of arrays,
+            one per embedding network) and "compress_info" (list of arrays
+            with table bounds).
+        """
+        variables = compress["@variables"]
+        self.compress_data = variables["compress_data"]
+        self.compress_info = variables["compress_info"]
+        self.compress = True
+
     @classmethod
     def update_sel(
         cls,

deepmd/dpmodel/descriptor/se_r.py

@@ -173,6 +173,7 @@ def __init__(
         self.type_map = type_map
         self.emask = PairExcludeMask(self.ntypes, self.exclude_types)
         self.env_protection = env_protection
+        self.compress = False
 
         in_dim = 1  # not considiering type embedding
         embeddings = NetworkCollection(
@@ -438,10 +439,10 @@ def call(
 
     def serialize(self) -> dict:
         """Serialize the descriptor to dict."""
-        return {
+        data = {
             "@class": "Descriptor",
             "type": "se_r",
-            "@version": 2,
+            "@version": 3 if self.compress else 2,
             "rcut": self.rcut,
             "rcut_smth": self.rcut_smth,
             "sel": self.sel,
@@ -464,24 +465,42 @@ def serialize(self) -> dict:
             },
             "type_map": self.type_map,
         }
+        if self.compress:
+            data["compress"] = {
+                "@variables": {
+                    "compress_data": [to_numpy_array(d) for d in self.compress_data],
+                    "compress_info": [to_numpy_array(i) for i in self.compress_info],
+                },
+            }
+        return data
 
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptSeR":
         """Deserialize from dict."""
         data = data.copy()
-        check_version_compatibility(data.pop("@version", 1), 2, 1)
+        check_version_compatibility(data.pop("@version", 1), 3, 1)
         data.pop("@class", None)
         data.pop("type", None)
         variables = data.pop("@variables")
         embeddings = data.pop("embeddings")
         env_mat = data.pop("env_mat")
+        compress = data.pop("compress", None)
         obj = cls(**data)
 
         obj["davg"] = variables["davg"]
         obj["dstd"] = variables["dstd"]
         obj.embeddings = NetworkCollection.deserialize(embeddings)
+        if compress is not None:
+            obj._load_compress_data(compress)
         return obj
 
+    def _load_compress_data(self, compress: dict) -> None:
+        """Load compression state from serialized data."""
+        variables = compress["@variables"]
+        self.compress_data = variables["compress_data"]
+        self.compress_info = variables["compress_info"]
+        self.compress = True
+
     @classmethod
     def update_sel(
         cls,