pd: support different label_dict in CINN (#4795)

pop unnecessary item when wrapping model with `jit.to_static`, so we can
support se_e2_a/dpa2/dpa3 without extra modification.

@njzjz can you give some suggestions for better code improvements? The
current approach of fetching data via `self.get_data` isn't very
concise.

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

- **Bug Fixes**
- Improved compatibility by dynamically matching label input
specifications to available label keys during model compilation when
CINN is enabled. This prevents errors caused by mismatched label keys at
runtime.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: HydrogenSulfate <[email protected]>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

deepmd/pd/train/training.py

@@ -607,6 +607,27 @@ def warm_up_linear(step, warmup_steps):
             )
 
             backend = "CINN" if CINN else None
+            # NOTE: This is a trick to decide the right input_spec for wrapper.forward
+            _, label_dict, _ = self.get_data(is_train=True)
+
+            # Define specification templates
+            spec_templates = {
+                "find_box": np.float32(1.0),
+                "find_coord": np.float32(1.0),
+                "find_numb_copy": np.float32(0.0),
+                "numb_copy": static.InputSpec([1, 1], "int64", name="numb_copy"),
+                "find_energy": np.float32(1.0),
+                "energy": static.InputSpec([1, 1], "float64", name="energy"),
+                "find_force": np.float32(1.0),
+                "force": static.InputSpec([1, -1, 3], "float64", name="force"),
+                "find_virial": np.float32(0.0),
+                "virial": static.InputSpec([1, 9], "float64", name="virial"),
+                "natoms": static.InputSpec([1, -1], "int32", name="natoms"),
+            }
+            # Build spec only for keys present in sample data
+            label_dict_spec = {
+                k: spec_templates[k] for k in label_dict.keys() if k in spec_templates
+            }
             self.wrapper.forward = jit.to_static(
                 backend=backend,
                 input_spec=[
@@ -615,19 +636,7 @@ def warm_up_linear(step, warmup_steps):
                     None,  # spin
                     static.InputSpec([1, 9], "float64", name="box"),  # box
                     static.InputSpec([], "float64", name="cur_lr"),  # cur_lr
-                    {
-                        "find_box": np.float32(1.0),
-                        "find_coord": np.float32(1.0),
-                        "find_numb_copy": np.float32(0.0),
-                        "numb_copy": static.InputSpec(
-                            [1, 1], "int64", name="numb_copy"
-                        ),
-                        "find_energy": np.float32(1.0),
-                        "energy": static.InputSpec([1, 1], "float64", name="energy"),
-                        "find_force": np.float32(1.0),
-                        "force": static.InputSpec([1, -1, 3], "float64", name="force"),
-                        "natoms": static.InputSpec([1, -1], "int32", name="natoms"),
-                    },  # label,
+                    label_dict_spec,  # label,
                     # None, # task_key
                     # False, # inference_only
                     # False, # do_atomic_virial