Merge branch 'devel' into devel-use_aparam_as_mask

.github/workflows/build_wheel.yml

@@ -148,6 +148,7 @@ jobs:
       - uses: actions/download-artifact@v4
         with:
           path: source/install/docker/dist
+          pattern: cibw-*-manylinux_x86_64-cu${{ matrix.cuda_version }}*
           merge-multiple: true
       - name: Log in to the Container registry
         uses: docker/login-action@v3
@@ -180,6 +181,7 @@ jobs:
       - uses: actions/download-artifact@v4
         with:
           path: dist/packages
+          pattern: cibw-*
           merge-multiple: true
       - uses: actions/setup-python@v5
         name: Install Python

.pre-commit-config.yaml

@@ -38,6 +38,14 @@ repos:
       - id: ruff-format
         exclude: ^source/3rdparty
         types_or: [python, pyi, jupyter]
+  - repo: https://github.com/pycqa/flake8
+    # flake8 cannot autofix
+    rev: "7.1.1"
+    hooks:
+      - id: flake8
+        additional_dependencies:
+          - torchfix==0.6.0
+          - flake8-pyproject==1.2.3
   # numpydoc
   - repo: https://github.com/Carreau/velin
     rev: 0.0.12

deepmd/dpmodel/array_api.py

@@ -61,7 +61,7 @@ def xp_take_along_axis(arr, indices, axis):
     else:
         indices = xp.reshape(indices, (0, 0))
 
-    offset = (xp.arange(indices.shape[0]) * m)[:, xp.newaxis]
+    offset = (xp.arange(indices.shape[0], dtype=indices.dtype) * m)[:, xp.newaxis]
     indices = xp.reshape(offset + indices, (-1,))
 
     out = xp.take(arr, indices)

deepmd/pt/entrypoints/main.py

@@ -105,8 +105,7 @@ def get_trainer(
     local_rank = os.environ.get("LOCAL_RANK")
     if local_rank is not None:
         local_rank = int(local_rank)
-        assert dist.is_nccl_available()
-        dist.init_process_group(backend="nccl")
+        dist.init_process_group(backend="cuda:nccl,cpu:gloo")
 
     def prepare_trainer_input_single(
         model_params_single, data_dict_single, rank=0, seed=None

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -74,8 +74,8 @@ def __init__(
         super().__init__(type_map, **kwargs)
         super().init_out_stat()
         self.tab_file = tab_file
-        self.rcut = rcut
-        self.tab = self._set_pairtab(tab_file, rcut)
+        self.rcut = float(rcut)
+        self.tab = self._set_pairtab(tab_file, self.rcut)
 
         self.type_map = type_map
         self.ntypes = len(type_map)

deepmd/pt/model/descriptor/repformer_layer.py

@@ -605,8 +605,8 @@ def __init__(
     ):
         super().__init__()
         self.epsilon = 1e-4  # protection of 1./nnei
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.ntypes = ntypes
         sel = [sel] if isinstance(sel, int) else sel
         self.nnei = sum(sel)

deepmd/pt/model/descriptor/repformers.py

@@ -193,8 +193,8 @@ def __init__(
             Random seed for parameter initialization.
         """
         super().__init__()
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.ntypes = ntypes
         self.nlayers = nlayers
         sel = [sel] if isinstance(sel, int) else sel

deepmd/pt/model/descriptor/se_a.py

@@ -395,8 +395,8 @@ def __init__(
         - axis_neuron: Number of columns of the sub-matrix of the embedding matrix.
         """
         super().__init__()
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.neuron = neuron
         self.filter_neuron = self.neuron
         self.axis_neuron = axis_neuron

deepmd/pt/model/descriptor/se_atten.py

@@ -149,8 +149,8 @@ def __init__(
         """
         super().__init__()
         del type
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.neuron = neuron
         self.filter_neuron = self.neuron
         self.axis_neuron = axis_neuron

deepmd/pt/model/descriptor/se_r.py

@@ -74,8 +74,8 @@ def __init__(
         **kwargs,
     ):
         super().__init__()
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.neuron = neuron
         self.filter_neuron = self.neuron
         self.set_davg_zero = set_davg_zero

deepmd/pt/model/descriptor/se_t.py

@@ -446,8 +446,8 @@ def __init__(
             Random seed for initializing the network parameters.
         """
         super().__init__()
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.neuron = neuron
         self.filter_neuron = self.neuron
         self.set_davg_zero = set_davg_zero

deepmd/pt/model/descriptor/se_t_tebd.py

@@ -512,8 +512,8 @@ def __init__(
         seed: Optional[Union[int, list[int]]] = None,
     ):
         super().__init__()
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.neuron = neuron
         self.filter_neuron = self.neuron
         self.tebd_dim = tebd_dim

deepmd/pt/utils/neighbor_stat.py

@@ -44,7 +44,7 @@ def __init__(
         mixed_types: bool,
     ) -> None:
         super().__init__()
-        self.rcut = rcut
+        self.rcut = float(rcut)
         self.ntypes = ntypes
         self.mixed_types = mixed_types
 

pyproject.toml

@@ -419,3 +419,10 @@ plugins = ["source.3rdparty.coverage_plugins.jit_plugin"]
 load-plugins = "deepmd_checker"
 disable = "all"
 enable = "E8001,E8002"
+
+[tool.flake8]
+select = [
+    "TOR0",
+    "TOR1",
+    "TOR2",
+]

source/checker/deepmd_checker.py

@@ -37,7 +37,7 @@ def visit_call(self, node):
         if (
             isinstance(node.func, Attribute)
             and isinstance(node.func.expr, Name)
-            and node.func.expr.name in {"np", "tf", "torch"}
+            and node.func.expr.name in {"np", "tf", "torch", "xp", "jnp"}
             and node.func.attrname
             in {
                 # https://pytorch.org/docs/stable/torch.html#creation-ops

source/tests/consistent/descriptor/test_dpa1.py

@@ -52,22 +52,22 @@
     (4,),  # tebd_dim
     ("concat", "strip"),  # tebd_input_mode
     (True,),  # resnet_dt
-    (True, False),  # type_one_side
+    (True,),  # type_one_side
     (20,),  # attn
     (0, 2),  # attn_layer
-    (True, False),  # attn_dotr
+    (True,),  # attn_dotr
     ([], [[0, 1]]),  # excluded_types
     (0.0,),  # env_protection
     (True, False),  # set_davg_zero
     (1.0,),  # scaling_factor
-    (True, False),  # normalize
+    (True,),  # normalize
     (None, 1.0),  # temperature
     (1e-5,),  # ln_eps
-    (True, False),  # smooth_type_embedding
+    (True,),  # smooth_type_embedding
     (True,),  # concat_output_tebd
     ("float64",),  # precision
     (True, False),  # use_econf_tebd
-    (False, True),  # use_tebd_bias
+    (False,),  # use_tebd_bias
 )
 class TestDPA1(CommonTest, DescriptorTest, unittest.TestCase):
     @property
@@ -127,11 +127,9 @@ def data(self) -> dict:
     def is_meaningless_zero_attention_layer_tests(
         self,
         attn_layer: int,
-        attn_dotr: bool,
-        normalize: bool,
         temperature: Optional[float],
     ) -> bool:
-        return attn_layer == 0 and (attn_dotr or normalize or temperature is not None)
+        return attn_layer == 0 and (temperature is not None)
 
     @property
     def skip_pt(self) -> bool:
@@ -158,8 +156,6 @@ def skip_pt(self) -> bool:
         ) = self.param
         return CommonTest.skip_pt or self.is_meaningless_zero_attention_layer_tests(
             attn_layer,
-            attn_dotr,
-            normalize,
             temperature,
         )
 
@@ -188,8 +184,6 @@ def skip_dp(self) -> bool:
         ) = self.param
         return CommonTest.skip_dp or self.is_meaningless_zero_attention_layer_tests(
             attn_layer,
-            attn_dotr,
-            normalize,
             temperature,
         )
 
@@ -218,8 +212,6 @@ def skip_jax(self) -> bool:
         ) = self.param
         return not INSTALLED_JAX or self.is_meaningless_zero_attention_layer_tests(
             attn_layer,
-            attn_dotr,
-            normalize,
             temperature,
         )
 
@@ -250,8 +242,6 @@ def skip_array_api_strict(self) -> bool:
             not INSTALLED_ARRAY_API_STRICT
             or self.is_meaningless_zero_attention_layer_tests(
                 attn_layer,
-                attn_dotr,
-                normalize,
                 temperature,
             )
         )
@@ -290,8 +280,6 @@ def skip_tf(self) -> bool:
             )
             or self.is_meaningless_zero_attention_layer_tests(
                 attn_layer,
-                attn_dotr,
-                normalize,
                 temperature,
             )
         )

source/tests/pt/model/test_jit.py

@@ -144,5 +144,27 @@ def tearDown(self):
         JITTest.tearDown(self)
 
 
+class TestEnergyModelDPA2IntRcut(unittest.TestCase, JITTest):
+    def setUp(self):
+        input_json = str(Path(__file__).parent / "water/se_atten.json")
+        with open(input_json) as f:
+            self.config = json.load(f)
+        data_file = [str(Path(__file__).parent / "water/data/data_0")]
+        self.config["training"]["training_data"]["systems"] = data_file
+        self.config["training"]["validation_data"]["systems"] = data_file
+        self.config["model"] = deepcopy(model_dpa2)
+        self.config["model"]["descriptor"]["repinit"]["rcut"] = int(
+            self.config["model"]["descriptor"]["repinit"]["rcut"]
+        )
+        self.config["model"]["descriptor"]["repinit"]["rcut_smth"] = int(
+            self.config["model"]["descriptor"]["repinit"]["rcut_smth"]
+        )
+        self.config["training"]["numb_steps"] = 10
+        self.config["training"]["save_freq"] = 10
+
+    def tearDown(self):
+        JITTest.tearDown(self)
+
+
 if __name__ == "__main__":
     unittest.main()