fix(jax): workaround for "xxTracer is not a valid JAX type" (#4776)

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

## Summary by CodeRabbit

- **Bug Fixes**
- Improved array handling in various components to ensure correct data
is passed during calculations, enhancing reliability and consistency in
model operations.
- **Style**
- Minor formatting adjustments for improved code readability, with no
impact on end-user functionality.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

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

Author: njzjz

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -293,7 +293,7 @@ def _pair_tabulated_inter(
 
         uu -= idx
         table_coef = self._extract_spline_coefficient(
-            i_type, j_type, idx, self.tab_data, nspline
+            i_type, j_type, idx, self.tab_data[...], nspline
         )
         table_coef = xp.reshape(table_coef, (nframes, nloc, nnei, 4))
         ener = self._calculate_ener(table_coef, uu)

deepmd/dpmodel/descriptor/dpa1.py

@@ -951,7 +951,11 @@ def call(
         xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext)
         # nf x nloc x nnei x 4
         dmatrix, diff, sw = self.env_mat.call(
-            coord_ext, atype_ext, nlist, self.mean, self.stddev
+            coord_ext,
+            atype_ext,
+            nlist,
+            self.mean[...],
+            self.stddev[...],
         )
         nf, nloc, nnei, _ = dmatrix.shape
         atype = atype_ext[:, :nloc]

deepmd/dpmodel/descriptor/repflows.py

@@ -472,7 +472,11 @@ def call(
         nlist = xp.where(exclude_mask, nlist, xp.full_like(nlist, -1))
         # nb x nloc x nnei x 4, nb x nloc x nnei x 3, nb x nloc x nnei x 1
         dmatrix, diff, sw = self.env_mat_edge.call(
-            coord_ext, atype_ext, nlist, self.mean, self.stddev
+            coord_ext,
+            atype_ext,
+            nlist,
+            self.mean[...],
+            self.stddev[...],
         )
         # nb x nloc x nnei
         nlist_mask = nlist != -1

deepmd/dpmodel/descriptor/repformers.py

@@ -441,7 +441,11 @@ def call(
         nlist = xp.where(exclude_mask, nlist, xp.full_like(nlist, -1))
         # nf x nloc x nnei x 4
         dmatrix, diff, sw = self.env_mat.call(
-            coord_ext, atype_ext, nlist, self.mean, self.stddev
+            coord_ext,
+            atype_ext,
+            nlist,
+            self.mean[...],
+            self.stddev[...],
         )
         nf, nloc, nnei, _ = dmatrix.shape
         # nf x nloc x nnei

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -591,7 +591,11 @@ def call(
         input_dtype = coord_ext.dtype
         # nf x nloc x nnei x 4
         rr, diff, ww = self.env_mat.call(
-            coord_ext, atype_ext, nlist, self.davg, self.dstd
+            coord_ext,
+            atype_ext,
+            nlist,
+            self.davg[...],
+            self.dstd[...],
         )
         nf, nloc, nnei, _ = rr.shape
         sec = self.sel_cumsum

deepmd/dpmodel/descriptor/se_r.py

@@ -373,7 +373,12 @@ def call(
         del mapping
         # nf x nloc x nnei x 1
         rr, diff, ww = self.env_mat.call(
-            coord_ext, atype_ext, nlist, self.davg, self.dstd, True
+            coord_ext,
+            atype_ext,
+            nlist,
+            self.davg[...],
+            self.dstd[...],
+            True,
         )
         nf, nloc, nnei, _ = rr.shape
         sec = self.sel_cumsum

deepmd/dpmodel/descriptor/se_t.py

@@ -349,7 +349,11 @@ def call(
         xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist)
         # nf x nloc x nnei x 4
         rr, diff, ww = self.env_mat.call(
-            coord_ext, atype_ext, nlist, self.davg, self.dstd
+            coord_ext,
+            atype_ext,
+            nlist,
+            self.davg[...],
+            self.dstd[...],
         )
         nf, nloc, nnei, _ = rr.shape
         sec = self.sel_cumsum

deepmd/dpmodel/descriptor/se_t_tebd.py

@@ -733,7 +733,11 @@ def call(
         xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext)
         # nf x nloc x nnei x 4
         dmatrix, diff, sw = self.env_mat.call(
-            coord_ext, atype_ext, nlist, self.mean, self.stddev
+            coord_ext,
+            atype_ext,
+            nlist,
+            self.mean[...],
+            self.stddev[...],
         )
         nf, nloc, nnei, _ = dmatrix.shape
         exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)

deepmd/dpmodel/fitting/general_fitting.py

@@ -410,7 +410,7 @@ def _call_common(
                     f"get an input fparam of dim {fparam.shape[-1]}, "
                     f"which is not consistent with {self.numb_fparam}."
                 )
-            fparam = (fparam - self.fparam_avg) * self.fparam_inv_std
+            fparam = (fparam - self.fparam_avg[...]) * self.fparam_inv_std[...]
             fparam = xp.tile(
                 xp.reshape(fparam, [nf, 1, self.numb_fparam]), (1, nloc, 1)
             )
@@ -432,7 +432,7 @@ def _call_common(
                     f"which is not consistent with {self.numb_aparam}."
                 )
             aparam = xp.reshape(aparam, [nf, nloc, self.numb_aparam])
-            aparam = (aparam - self.aparam_avg) * self.aparam_inv_std
+            aparam = (aparam - self.aparam_avg[...]) * self.aparam_inv_std[...]
             xx = xp.concat(
                 [xx, aparam],
                 axis=-1,
@@ -445,7 +445,9 @@ def _call_common(
 
         if self.dim_case_embd > 0:
             assert self.case_embd is not None
-            case_embd = xp.tile(xp.reshape(self.case_embd, [1, 1, -1]), [nf, nloc, 1])
+            case_embd = xp.tile(
+                xp.reshape(self.case_embd[...], [1, 1, -1]), [nf, nloc, 1]
+            )
             xx = xp.concat(
                 [xx, case_embd],
                 axis=-1,
@@ -482,7 +484,9 @@ def _call_common(
                 outs -= self.nets[()](xx_zeros)
         outs += xp.reshape(
             xp.take(
-                xp.astype(self.bias_atom_e, outs.dtype), xp.reshape(atype, [-1]), axis=0
+                xp.astype(self.bias_atom_e[...], outs.dtype),
+                xp.reshape(atype, [-1]),
+                axis=0,
             ),
             [nf, nloc, net_dim_out],
         )

deepmd/dpmodel/utils/exclude_mask.py

@@ -53,7 +53,8 @@ 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), (nf, natom)
+            xp.take(self.type_mask[...], xp.reshape(atype, [-1]), axis=0),
+            (nf, natom),
         )
 
 
@@ -131,7 +132,8 @@ def build_type_exclude_mask(
         # nf x (nloc x nnei)
         type_ij = xp.reshape(type_ij, (nf, nloc * nnei))
         mask = xp.reshape(
-            xp.take(self.type_mask, xp.reshape(type_ij, (-1,))), (nf, nloc, nnei)
+            xp.take(self.type_mask[...], xp.reshape(type_ij, (-1,))),
+            (nf, nloc, nnei),
         )
         return mask