diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index 0c35320e7f..9d7739d5c8 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/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) diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 3b25e63fb5..51c56e9681 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/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] diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index 926b500645..4b9922596c 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/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 diff --git a/deepmd/dpmodel/descriptor/repformers.py b/deepmd/dpmodel/descriptor/repformers.py index a917f7a227..3d02054350 100644 --- a/deepmd/dpmodel/descriptor/repformers.py +++ b/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 diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index d38955b98f..bd72d936e3 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/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 diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index 01df14394e..5b2931b23f 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/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 diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py index 3540eae53c..fb30f04961 100644 --- a/deepmd/dpmodel/descriptor/se_t.py +++ b/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 diff --git a/deepmd/dpmodel/descriptor/se_t_tebd.py b/deepmd/dpmodel/descriptor/se_t_tebd.py index c3ee41fe00..ff26024aad 100644 --- a/deepmd/dpmodel/descriptor/se_t_tebd.py +++ b/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) diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 7342663141..cd0d4e72d4 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/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], ) diff --git a/deepmd/dpmodel/utils/exclude_mask.py b/deepmd/dpmodel/utils/exclude_mask.py index 372eba133e..f390bbc7c1 100644 --- a/deepmd/dpmodel/utils/exclude_mask.py +++ b/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 diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py index 6a3c6d8081..ed5e90b81b 100644 --- a/deepmd/dpmodel/utils/network.py +++ b/deepmd/dpmodel/utils/network.py @@ -259,9 +259,9 @@ def call(self, x: np.ndarray) -> np.ndarray: xp = array_api_compat.array_namespace(x) fn = get_activation_fn(self.activation_function) y = ( - xp.matmul(x, self.w) + self.b + xp.matmul(x, self.w[...]) + self.b[...] if self.b is not None - else xp.matmul(x, self.w) + else xp.matmul(x, self.w[...]) ) if y.dtype != x.dtype: # workaround for bfloat16