fix jax reshape list

Author: caic99

deepmd/dpmodel/descriptor/dpa1.py

@@ -520,7 +520,7 @@ def call(
         type_embedding = self.type_embedding.call()
         # nf x nall x tebd_dim
         atype_embd_ext = xp.reshape(
-            xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0),
+            xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
             (nf, nall, self.tebd_dim),
         )
         # nfnl x tebd_dim

deepmd/dpmodel/descriptor/dpa2.py

@@ -841,7 +841,7 @@ def call(
         type_embedding = self.type_embedding.call()
         # repinit
         g1_ext = xp.reshape(
-            xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0),
+            xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
             (nframes, nall, self.tebd_dim),
         )
         g1_inp = g1_ext[:, :nloc, :]

deepmd/dpmodel/descriptor/dpa3.py

@@ -562,12 +562,12 @@ def call(
         type_embedding = self.type_embedding.call()
         if self.use_loc_mapping:
             node_ebd_ext = xp.reshape(
-                xp.take(type_embedding, xp.reshape(atype_ext[:, :nloc], [-1]), axis=0),
+                xp.take(type_embedding, xp.reshape(atype_ext[:, :nloc], (-1,)), axis=0),
                 (nframes, nloc, self.tebd_dim),
             )
         else:
             node_ebd_ext = xp.reshape(
-                xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0),
+                xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
                 (nframes, nall, self.tebd_dim),
             )
         node_ebd_inp = node_ebd_ext[:, :nloc, :]

deepmd/dpmodel/descriptor/se_t_tebd.py

@@ -358,7 +358,7 @@ def call(
         type_embedding = self.type_embedding.call()
         # nf x nall x tebd_dim
         atype_embd_ext = xp.reshape(
-            xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0),
+            xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
             (nf, nall, self.tebd_dim),
         )
         # nfnl x tebd_dim

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -289,7 +289,7 @@ def call(
         ]
         # out = out * self.scale[atype, ...]
         scale_atype = xp.reshape(
-            xp.take(xp.astype(self.scale, out.dtype), xp.reshape(atype, [-1]), axis=0),
+            xp.take(xp.astype(self.scale, out.dtype), xp.reshape(atype, (-1,)), axis=0),
             (*atype.shape, 1),
         )
         out = out * scale_atype
@@ -315,7 +315,7 @@ def call(
             bias = xp.reshape(
                 xp.take(
                     xp.astype(self.constant_matrix, out.dtype),
-                    xp.reshape(atype, [-1]),
+                    xp.reshape(atype, (-1,)),
                     axis=0,
                 ),
                 (nframes, nloc),

deepmd/dpmodel/loss/ener.py

@@ -132,18 +132,18 @@ def call(
             atom_ener_coeff = xp.reshape(atom_ener_coeff, xp.shape(atom_ener))
             energy = xp.sum(atom_ener_coeff * atom_ener, 1)
         if self.has_f or self.has_pf or self.relative_f or self.has_gf:
-            force_reshape = xp.reshape(force, [-1])
-            force_hat_reshape = xp.reshape(force_hat, [-1])
+            force_reshape = xp.reshape(force, (-1,))
+            force_hat_reshape = xp.reshape(force_hat, (-1,))
             diff_f = force_hat_reshape - force_reshape
         else:
             diff_f = None
 
         if self.relative_f is not None:
-            force_hat_3 = xp.reshape(force_hat, [-1, 3])
-            norm_f = xp.reshape(xp.norm(force_hat_3, axis=1), [-1, 1]) + self.relative_f
-            diff_f_3 = xp.reshape(diff_f, [-1, 3])
+            force_hat_3 = xp.reshape(force_hat, (-1, 3))
+            norm_f = xp.reshape(xp.norm(force_hat_3, axis=1), (-1, 1)) + self.relative_f
+            diff_f_3 = xp.reshape(diff_f, (-1, 3))
             diff_f_3 = diff_f_3 / norm_f
-            diff_f = xp.reshape(diff_f_3, [-1])
+            diff_f = xp.reshape(diff_f_3, (-1,))
 
         atom_norm = 1.0 / natoms
         atom_norm_ener = 1.0 / natoms
@@ -184,15 +184,15 @@ def call(
                 loss += pref_f * l2_force_loss
             else:
                 l_huber_loss = custom_huber_loss(
-                    xp.reshape(force, [-1]),
-                    xp.reshape(force_hat, [-1]),
+                    xp.reshape(force, (-1,)),
+                    xp.reshape(force_hat, (-1,)),
                     delta=self.huber_delta,
                 )
                 loss += pref_f * l_huber_loss
             more_loss["rmse_f"] = self.display_if_exist(l2_force_loss, find_force)
         if self.has_v:
-            virial_reshape = xp.reshape(virial, [-1])
-            virial_hat_reshape = xp.reshape(virial_hat, [-1])
+            virial_reshape = xp.reshape(virial, (-1,))
+            virial_hat_reshape = xp.reshape(virial_hat, (-1,))
             l2_virial_loss = xp.mean(
                 xp.square(virial_hat_reshape - virial_reshape),
             )
@@ -207,8 +207,8 @@ def call(
                 loss += pref_v * l_huber_loss
             more_loss["rmse_v"] = self.display_if_exist(l2_virial_loss, find_virial)
         if self.has_ae:
-            atom_ener_reshape = xp.reshape(atom_ener, [-1])
-            atom_ener_hat_reshape = xp.reshape(atom_ener_hat, [-1])
+            atom_ener_reshape = xp.reshape(atom_ener, (-1,))
+            atom_ener_hat_reshape = xp.reshape(atom_ener_hat, (-1,))
             l2_atom_ener_loss = xp.mean(
                 xp.square(atom_ener_hat_reshape - atom_ener_reshape),
             )
@@ -225,7 +225,7 @@ def call(
                 l2_atom_ener_loss, find_atom_ener
             )
         if self.has_pf:
-            atom_pref_reshape = xp.reshape(atom_pref, [-1])
+            atom_pref_reshape = xp.reshape(atom_pref, (-1,))
             l2_pref_force_loss = xp.mean(
                 xp.multiply(xp.square(diff_f), atom_pref_reshape),
             )
@@ -236,10 +236,10 @@ def call(
         if self.has_gf:
             find_drdq = label_dict["find_drdq"]
             drdq = label_dict["drdq"]
-            force_reshape_nframes = xp.reshape(force, [-1, natoms[0] * 3])
-            force_hat_reshape_nframes = xp.reshape(force_hat, [-1, natoms[0] * 3])
+            force_reshape_nframes = xp.reshape(force, (-1, natoms[0] * 3))
+            force_hat_reshape_nframes = xp.reshape(force_hat, (-1, natoms[0] * 3))
             drdq_reshape = xp.reshape(
-                drdq, [-1, natoms[0] * 3, self.numb_generalized_coord]
+                drdq, (-1, natoms[0] * 3, self.numb_generalized_coord)
             )
             gen_force_hat = xp.einsum(
                 "bij,bi->bj", drdq_reshape, force_hat_reshape_nframes

deepmd/dpmodel/model/transform_output.py

@@ -100,7 +100,7 @@ def communicate_extended_output(
             if vdef.r_differentiable:
                 if model_ret[kk_derv_r] is not None:
                     derv_r_ext_dims = list(vdef.shape) + [3]  # noqa:RUF005
-                    mapping = xp.reshape(mapping, (mldims + [1] * len(derv_r_ext_dims)))
+                    mapping = xp.reshape(mapping, tuple(mldims + [1] * len(derv_r_ext_dims)))
                     mapping = xp.tile(mapping, [1] * len(mldims) + derv_r_ext_dims)
                     force = xp.zeros(vldims + derv_r_ext_dims, dtype=vv.dtype)
                     force = xp_scatter_sum(

deepmd/dpmodel/utils/env_mat_stat.py

@@ -189,7 +189,7 @@ def iter(
             for type_i in range(self.descriptor.get_ntypes()):
                 dd = env_mat[type_idx[type_i, ...]]
                 dd = xp.reshape(
-                    dd, [-1, self.last_dim]
+                    dd, (-1, self.last_dim)
                 )  # typen_atoms * unmasked_nnei, 4
                 env_mats = {}
                 env_mats[f"r_{type_i}"] = dd[:, :1]

deepmd/dpmodel/utils/exclude_mask.py

@@ -53,7 +53,7 @@ def build_type_exclude_mask(
         xp = array_api_compat.array_namespace(atype)
         nf, natom = atype.shape
         return xp.reshape(
-            xp.take(self.type_mask[...], xp.reshape(atype, [-1]), axis=0),
+            xp.take(self.type_mask[...], xp.reshape(atype, (-1,)), axis=0),
             (nf, natom),
         )
 

deepmd/dpmodel/utils/neighbor_stat.py

@@ -82,8 +82,8 @@ def call(
         nall = coord1.shape[1] // 3
         coord0 = coord1[:, : nloc * 3]
         diff = (
-            xp.reshape(coord1, [nframes, -1, 3])[:, None, :, :]
-            - xp.reshape(coord0, [nframes, -1, 3])[:, :, None, :]
+            xp.reshape(coord1, (nframes, -1, 3))[:, None, :, :]
+            - xp.reshape(coord0, (nframes, -1, 3))[:, :, None, :]
         )
         assert list(diff.shape) == [nframes, nloc, nall, 3]
         # remove the diagonal elements