fix(pt): fix type annotations for dummy compress op; improve docs

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/pt/model/descriptor/se_a.py

@@ -73,11 +73,11 @@
 if not hasattr(torch.ops.deepmd, "tabulate_fusion_se_a"):
 
     def tabulate_fusion_se_a(
-        argument0,
-        argument1,
-        argument2,
-        argument3,
-        argument4,
+        argument0: torch.Tensor,
+        argument1: torch.Tensor,
+        argument2: torch.Tensor,
+        argument3: torch.Tensor,
+        argument4: int,
     ) -> list[torch.Tensor]:
         raise NotImplementedError(
             "tabulate_fusion_se_a is not available since customized PyTorch OP library is not built when freezing the model. "

deepmd/pt/model/descriptor/se_atten.py

@@ -52,13 +52,13 @@
 if not hasattr(torch.ops.deepmd, "tabulate_fusion_se_atten"):
 
     def tabulate_fusion_se_atten(
-        argument0,
-        argument1,
-        argument2,
-        argument3,
-        argument4,
-        argument5,
-        argument6,
+        argument0: torch.Tensor,
+        argument1: torch.Tensor,
+        argument2: torch.Tensor,
+        argument3: torch.Tensor,
+        argument4: torch.Tensor,
+        argument5: int,
+        argument6: bool,
     ) -> list[torch.Tensor]:
         raise NotImplementedError(
             "tabulate_fusion_se_atten is not available since customized PyTorch OP library is not built when freezing the model. "

deepmd/pt/model/descriptor/se_r.py

@@ -62,10 +62,10 @@
 if not hasattr(torch.ops.deepmd, "tabulate_fusion_se_r"):
 
     def tabulate_fusion_se_r(
-        argument0,
-        argument1,
-        argument2,
-        argument3,
+        argument0: torch.Tensor,
+        argument1: torch.Tensor,
+        argument2: torch.Tensor,
+        argument3: int,
     ) -> list[torch.Tensor]:
         raise NotImplementedError(
             "tabulate_fusion_se_r is not available since customized PyTorch OP library is not built when freezing the model. "

deepmd/pt/model/descriptor/se_t.py

@@ -73,11 +73,11 @@
 if not hasattr(torch.ops.deepmd, "tabulate_fusion_se_t"):
 
     def tabulate_fusion_se_t(
-        argument0,
-        argument1,
-        argument2,
-        argument3,
-        argument4,
+        argument0: torch.Tensor,
+        argument1: torch.Tensor,
+        argument2: torch.Tensor,
+        argument3: torch.Tensor,
+        argument4: int,
     ) -> list[torch.Tensor]:
         raise NotImplementedError(
             "tabulate_fusion_se_t is not available since customized PyTorch OP library is not built when freezing the model. "

doc/freeze/compress.md

@@ -124,3 +124,9 @@ Notice: Model compression for the `se_atten_v2` descriptor is exclusively design
 - relu6
 - softplus
 - sigmoid
+
+## Requirements of installation {{ pytorch_icon }}
+
+When compressing models in the PyTorch backend, the customized OP library for the Python interface must be installed when [freezing the model](../freeze/freeze.md).
+
+The customized OP library for the Python interface can be installed by setting environment variable {envvar}`DP_ENABLE_PYTORCH` to `1` during [installation](../install/install-from-source.md).