fix typos

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>

Author: njzjz

deepmd/dpmodel/descriptor/dpa1.py

@@ -896,10 +896,10 @@ def compute_input_stats(
         env_mat_stat.load_or_compute_stats(sampled, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
-        xp = array_api_compat.array_namespace(self.dstd)
+        xp = array_api_compat.array_namespace(self.stddev)
         if not self.set_davg_zero:
-            self.davg = xp.asarray(mean, dtype=self.davg.dtype, copy=True)
-        self.dstd = xp.asarray(stddev, dtype=self.dstd.dtype, copy=True)
+            self.mean = xp.asarray(mean, dtype=self.mean.dtype, copy=True)
+        self.stddev = xp.asarray(stddev, dtype=self.stddev.dtype, copy=True)
 
     def get_stats(self) -> dict[str, StatItem]:
         """Get the statistics of the descriptor."""

deepmd/dpmodel/descriptor/repflows.py

@@ -385,10 +385,10 @@ def compute_input_stats(
         env_mat_stat.load_or_compute_stats(sampled, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
-        xp = array_api_compat.array_namespace(self.dstd)
+        xp = array_api_compat.array_namespace(self.stddev)
         if not self.set_davg_zero:
-            self.davg = xp.asarray(mean, dtype=self.davg.dtype, copy=True)
-        self.dstd = xp.asarray(stddev, dtype=self.dstd.dtype, copy=True)
+            self.mean = xp.asarray(mean, dtype=self.mean.dtype, copy=True)
+        self.stddev = xp.asarray(stddev, dtype=self.stddev.dtype, copy=True)
 
     def get_stats(self) -> dict[str, StatItem]:
         """Get the statistics of the descriptor."""

deepmd/dpmodel/descriptor/repformers.py

@@ -406,10 +406,10 @@ def compute_input_stats(
         env_mat_stat.load_or_compute_stats(sampled, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
-        xp = array_api_compat.array_namespace(self.dstd)
+        xp = array_api_compat.array_namespace(self.stddev)
         if not self.set_davg_zero:
-            self.davg = xp.asarray(mean, dtype=self.davg.dtype, copy=True)
-        self.dstd = xp.asarray(stddev, dtype=self.dstd.dtype, copy=True)
+            self.mean = xp.asarray(mean, dtype=self.mean.dtype, copy=True)
+        self.stddev = xp.asarray(stddev, dtype=self.stddev.dtype, copy=True)
 
     def get_stats(self) -> dict[str, StatItem]:
         """Get the statistics of the descriptor."""

deepmd/dpmodel/descriptor/se_t_tebd.py

@@ -682,10 +682,10 @@ def compute_input_stats(
         env_mat_stat.load_or_compute_stats(sampled, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
-        xp = array_api_compat.array_namespace(self.dstd)
+        xp = array_api_compat.array_namespace(self.stddev)
         if not self.set_davg_zero:
-            self.davg = xp.asarray(mean, dtype=self.davg.dtype, copy=True)
-        self.dstd = xp.asarray(stddev, dtype=self.dstd.dtype, copy=True)
+            self.mean = xp.asarray(mean, dtype=self.mean.dtype, copy=True)
+        self.stddev = xp.asarray(stddev, dtype=self.stddev.dtype, copy=True)
 
     def get_stats(self) -> dict[str, StatItem]:
         """Get the statistics of the descriptor."""

deepmd/dpmodel/utils/env_mat_stat.py

@@ -185,7 +185,7 @@ def iter(
                     (1, coord.shape[0] * coord.shape[1], -1),
                 )
                 # shape: (ntypes, nloc, nnei)
-                type_idx = xp.logical_and(type_idx.unsqueeze(-1), exclude_mask)
+                type_idx = xp.logical_and(type_idx[..., None], exclude_mask)
             for type_i in range(self.descriptor.get_ntypes()):
                 dd = env_mat[type_idx[type_i, ...]]
                 dd = xp.reshape(

source/tests/consistent/descriptor/test_se_e2_a.py

@@ -530,9 +530,9 @@ def eval_pd(self, pd_obj: Any) -> Any:
             [
                 {
                     "r0": None,
-                    "coord": paddle.to_tensor(self.coords)
-                    .reshape(-1, self.natoms[0], 3)
-                    .to(PD_DEVICE),
+                    "coord": paddle.to_tensor(
+                        self.coords.reshape(-1, self.natoms[0], 3)
+                    ).to(PD_DEVICE),
                     "atype": paddle.to_tensor(self.atype.reshape(1, -1)).to(PD_DEVICE),
                     "box": paddle.to_tensor(self.box.reshape(1, 3, 3)).to(PD_DEVICE),
                     "natoms": self.natoms[0],