fix uts

Author: iProzd

deepmd/dpmodel/descriptor/repflows.py

@@ -1294,7 +1294,7 @@ def call(
         )
         nb, nloc, nnei = nlist.shape
         nall = node_ebd_ext.shape[1]
-        n_edge = int(nlist_mask.sum().item())
+        n_edge = int(xp.sum(nlist_mask.astype(xp.int32)))
         node_ebd = node_ebd_ext[:, :nloc, :]
         assert (nb, nloc) == node_ebd.shape[:2]
         if not self.use_dynamic_sel:

deepmd/dpmodel/utils/network.py

@@ -1044,7 +1044,7 @@ def get_graph_index(
         a_nlist_mask[:, :, :, None], a_nlist_mask[:, :, None, :]
     )
 
-    n_edge = int(xp.asarray(xp.sum(nlist_mask)))
+    n_edge = int(xp.sum(nlist_mask.astype(xp.int32)))
 
     # following: get n2e_index, n_ext2e_index, n2a_index, eij2a_index, eik2a_index
 

deepmd/pd/model/descriptor/dpa3.py

@@ -150,6 +150,8 @@ def init_subclass_params(sub_data, sub_class):
             fix_stat_std=self.repflow_args.fix_stat_std,
             optim_update=self.repflow_args.optim_update,
             smooth_edge_update=self.repflow_args.smooth_edge_update,
+            use_dynamic_sel=self.repflow_args.use_dynamic_sel,
+            sel_reduce_factor=self.repflow_args.sel_reduce_factor,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,

deepmd/pd/model/descriptor/repflow_layer.py

@@ -52,6 +52,8 @@ def __init__(
         axis_neuron: int = 4,
         update_angle: bool = True,
         optim_update: bool = True,
+        use_dynamic_sel: bool = False,
+        sel_reduce_factor: float = 10.0,
         smooth_edge_update: bool = False,
         activation_function: str = "silu",
         update_style: str = "res_residual",
@@ -98,6 +100,10 @@ def __init__(
         self.prec = PRECISION_DICT[precision]
         self.optim_update = optim_update
         self.smooth_edge_update = smooth_edge_update
+        self.use_dynamic_sel = use_dynamic_sel
+        self.sel_reduce_factor = sel_reduce_factor
+        self.dynamic_e_sel = self.nnei / self.sel_reduce_factor
+        self.dynamic_a_sel = self.a_sel / self.sel_reduce_factor
 
         assert update_residual_init in [
             "norm",
@@ -812,7 +818,7 @@ def serialize(self) -> dict:
         """
         data = {
             "@class": "RepFlowLayer",
-            "@version": 1,
+            "@version": 2,
             "e_rcut": self.e_rcut,
             "e_rcut_smth": self.e_rcut_smth,
             "e_sel": self.e_sel,
@@ -836,6 +842,8 @@ def serialize(self) -> dict:
             "precision": self.precision,
             "optim_update": self.optim_update,
             "smooth_edge_update": self.smooth_edge_update,
+            "use_dynamic_sel": self.use_dynamic_sel,
+            "sel_reduce_factor": self.sel_reduce_factor,
             "node_self_mlp": self.node_self_mlp.serialize(),
             "node_sym_linear": self.node_sym_linear.serialize(),
             "node_edge_linear": self.node_edge_linear.serialize(),
@@ -878,7 +886,7 @@ def deserialize(cls, data: dict) -> "RepFlowLayer":
             The dict to deserialize from.
         """
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 1, 1)
+        check_version_compatibility(data.pop("@version"), 2, 1)
         data.pop("@class")
         update_angle = data["update_angle"]
         a_compress_rate = data["a_compress_rate"]

deepmd/pd/model/descriptor/repflows.py

@@ -159,6 +159,8 @@ def __init__(
         precision: str = "float64",
         fix_stat_std: float = 0.3,
         smooth_edge_update: bool = False,
+        use_dynamic_sel: bool = False,
+        sel_reduce_factor: float = 10.0,
         optim_update: bool = True,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
@@ -191,6 +193,9 @@ def __init__(
         self.a_compress_use_split = a_compress_use_split
         self.optim_update = optim_update
         self.smooth_edge_update = smooth_edge_update
+        self.use_dynamic_sel = use_dynamic_sel  # not supported yet
+        self.sel_reduce_factor = sel_reduce_factor
+        assert not self.use_dynamic_sel, "Dynamic selection is not supported yet."
 
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -243,6 +248,8 @@ def __init__(
                     update_residual_init=self.update_residual_init,
                     precision=precision,
                     optim_update=self.optim_update,
+                    use_dynamic_sel=self.use_dynamic_sel,
+                    sel_reduce_factor=self.sel_reduce_factor,
                     smooth_edge_update=self.smooth_edge_update,
                     seed=child_seed(child_seed(seed, 1), ii),
                 )

source/tests/universal/dpmodel/descriptor/test_descriptor.py

@@ -509,6 +509,7 @@ def DescriptorParamDPA3(
                 "n_multi_edge_message": n_multi_edge_message,
                 "axis_neuron": 2,
                 "use_dynamic_sel": use_dynamic_sel,
+                "sel_reduce_factor": 1.0,
                 "update_angle": update_angle,
                 "update_style": update_style,
                 "update_residual": update_residual,