From 7690428821be9cacd8bcad6809159b523795a317 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 8 Jan 2026 01:48:45 +0800 Subject: [PATCH 01/12] style: enable B905 to prevent issues with zip Signed-off-by: Jinzhe Zeng --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 5b089c4558..1e156f77c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -382,6 +382,7 @@ select = [ "TCH", # flake8-type-checking "PYI", # flake8-pyi "ANN", # type annotations + "B905", # zip-without-explicit-strict ] ignore = [ From f4cf02b934a86a18a30a6f52ffa26a4a188c0ad5 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 8 Jan 2026 01:54:54 +0800 Subject: [PATCH 02/12] run ruff --- deepmd/calculator.py | 2 +- .../atomic_model/linear_atomic_model.py | 6 ++--- deepmd/dpmodel/descriptor/hybrid.py | 4 ++-- deepmd/dpmodel/infer/deep_eval.py | 2 +- deepmd/dpmodel/utils/nlist.py | 2 +- deepmd/jax/infer/deep_eval.py | 2 +- deepmd/pd/infer/deep_eval.py | 2 +- deepmd/pd/model/descriptor/se_a.py | 2 +- deepmd/pd/model/model/transform_output.py | 2 +- deepmd/pd/utils/dataloader.py | 2 +- deepmd/pd/utils/nlist.py | 2 +- deepmd/pt/infer/deep_eval.py | 2 +- .../model/atomic_model/linear_atomic_model.py | 4 ++-- deepmd/pt/model/descriptor/hybrid.py | 2 +- deepmd/pt/model/descriptor/se_a.py | 2 +- deepmd/pt/model/descriptor/se_r.py | 2 +- deepmd/pt/model/descriptor/se_t.py | 2 +- deepmd/pt/model/model/transform_output.py | 2 +- deepmd/pt/optimizer/LKF.py | 2 +- deepmd/pt/train/training.py | 2 +- deepmd/pt/utils/dataloader.py | 2 +- deepmd/pt/utils/nlist.py | 2 +- deepmd/tf/descriptor/loc_frame.py | 2 +- deepmd/tf/descriptor/se_a.py | 4 ++-- deepmd/tf/descriptor/se_a_ef.py | 2 +- deepmd/tf/descriptor/se_atten.py | 4 ++-- deepmd/tf/descriptor/se_r.py | 2 +- deepmd/tf/descriptor/se_t.py | 2 +- deepmd/tf/infer/deep_eval.py | 2 +- deepmd/tf/nvnmd/entrypoints/mapt.py | 10 ++++---- deepmd/tf/train/trainer.py | 2 +- deepmd/utils/batch_size.py | 2 +- deepmd/utils/data_system.py | 2 +- deepmd/utils/model_branch_dict.py | 4 ++-- deepmd/utils/update_sel.py | 2 +- .../dpmodel/test_pairtab_atomic_model.py | 2 +- source/tests/consistent/common.py | 24 +++++++++---------- .../tests/consistent/fitting/test_dipole.py | 2 +- source/tests/consistent/io/test_io.py | 4 ++-- source/tests/pd/model/test_nlist.py | 2 +- source/tests/pd/model/test_se_e2_a.py | 2 +- source/tests/pd/test_utils.py | 2 +- source/tests/pt/model/test_descriptor_se_r.py | 2 +- source/tests/pt/model/test_nlist.py | 2 +- .../pt/model/test_pairtab_atomic_model.py | 2 +- source/tests/pt/model/test_se_e2_a.py | 2 +- source/tests/pt/model/test_se_t.py | 2 +- source/tests/pt/test_utils.py | 2 +- source/tests/tf/test_linear_model.py | 2 +- source/tests/tf/test_model_devi.py | 4 ++-- source/tests/tf/test_model_devi_mix.py | 2 +- source/tests/tf/test_nvnmd_entrypoints.py | 4 ++-- 52 files changed, 77 insertions(+), 77 deletions(-) diff --git a/deepmd/calculator.py b/deepmd/calculator.py index dca2343ce0..1b40cc09d6 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -103,7 +103,7 @@ def __init__( self.type_dict = type_dict else: self.type_dict = dict( - zip(self.dp.get_type_map(), range(self.dp.get_ntypes())) + zip(self.dp.get_type_map(), range(self.dp.get_ntypes()), strict=False) ) def calculate( diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 3fa1119418..743ca1cf39 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -155,7 +155,7 @@ def get_model_sels(self) -> list[int | list[int]]: def _sort_rcuts_sels(self) -> tuple[tuple[Array, Array], list[int]]: # sort the pair of rcut and sels in ascending order, first based on sel, then on rcut. zipped = sorted( - zip(self.get_model_rcuts(), self.get_model_nsels()), + zip(self.get_model_rcuts(), self.get_model_nsels(), strict=False), key=lambda x: (x[1], x[0]), ) return [p[0] for p in zipped], [p[1] for p in zipped] @@ -235,12 +235,12 @@ def forward_atomic( ) raw_nlists = [ nlists[get_multiple_nlist_key(rcut, sel)] - for rcut, sel in zip(self.get_model_rcuts(), self.get_model_nsels()) + for rcut, sel in zip(self.get_model_rcuts(), self.get_model_nsels(), strict=False) ] nlists_ = [ nl if mt else nlist_distinguish_types(nl, extended_atype, sel) for mt, nl, sel in zip( - self.mixed_types_list, raw_nlists, self.get_model_sels() + self.mixed_types_list, raw_nlists, self.get_model_sels(), strict=False ) ] ener_list = [] diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 3682e17c4d..95d4716d5c 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -101,7 +101,7 @@ def __init__( start_idx = np.cumsum(np.pad(hybrid_sel, (1, 0), "constant"))[:-1] end_idx = start_idx + np.array(sub_sel) cut_idx = np.concatenate( - [range(ss, ee) for ss, ee in zip(start_idx, end_idx)] + [range(ss, ee) for ss, ee in zip(start_idx, end_idx, strict=False)] ) nlist_cut_idx.append(cut_idx) self.nlist_cut_idx = nlist_cut_idx @@ -310,7 +310,7 @@ def call( ) else: nl_distinguish_types = None - for descrpt, nci in zip(self.descrpt_list, self.nlist_cut_idx): + for descrpt, nci in zip(self.descrpt_list, self.nlist_cut_idx, strict=False): # cut the nlist to the correct length if self.mixed_types() == descrpt.mixed_types(): nl = xp.take(nlist, nci, axis=2) diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index 3bd0f435e8..4526680fc4 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -228,7 +228,7 @@ def eval( return dict( zip( [x.name for x in request_defs], - out, + out, strict=False, ) ) diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index a43cf46403..1e361589c4 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -234,7 +234,7 @@ def build_multiple_neighbor_list( rr = xp.where(nlist_mask, xp.full_like(rr, float("inf")), rr) nlist0 = nlist ret = {} - for rc, ns in zip(rcuts[::-1], nsels[::-1]): + for rc, ns in zip(rcuts[::-1], nsels[::-1], strict=False): tnlist_1 = nlist0[:, :, :ns] tnlist_1 = xp.where(rr[:, :, :ns] > rc, xp.full_like(tnlist_1, -1), tnlist_1) ret[get_multiple_nlist_key(rc, ns)] = tnlist_1 diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 1e29ee1c78..1ae9662080 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -248,7 +248,7 @@ def eval( return dict( zip( [x.name for x in request_defs], - out, + out, strict=False, ) ) diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index 67c435ab3f..d40f91cb8d 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -384,7 +384,7 @@ def eval( return dict( zip( [x.name for x in request_defs], - out, + out, strict=False, ) ) diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py index 17ef5d67c4..dbed543041 100644 --- a/deepmd/pd/model/descriptor/se_a.py +++ b/deepmd/pd/model/descriptor/se_a.py @@ -763,7 +763,7 @@ def forward( # nfnl x nnei exclude_mask = self.emask(nlist, extended_atype).reshape([nfnl, self.nnei]) for embedding_idx, (ll, compress_data_ii, compress_info_ii) in enumerate( - zip(self.filter_layers.networks, self.compress_data, self.compress_info) + zip(self.filter_layers.networks, self.compress_data, self.compress_info, strict=False) ): if self.type_one_side: ii = embedding_idx diff --git a/deepmd/pd/model/model/transform_output.py b/deepmd/pd/model/model/transform_output.py index 47004265c7..692e01a40c 100644 --- a/deepmd/pd/model/model/transform_output.py +++ b/deepmd/pd/model/model/transform_output.py @@ -122,7 +122,7 @@ def take_deriv( split_vv1 = paddle.split(vv1, [1] * size, axis=-1) split_svv1 = paddle.split(svv1, [1] * size, axis=-1) split_ff, split_avir = [], [] - for vvi, svvi in zip(split_vv1, split_svv1): + for vvi, svvi in zip(split_vv1, split_svv1, strict=False): # nf x nloc x 3, nf x nloc x 9 ffi, aviri = task_deriv_one( vvi, diff --git a/deepmd/pd/utils/dataloader.py b/deepmd/pd/utils/dataloader.py index 0cb8adbc63..cf36ea2d05 100644 --- a/deepmd/pd/utils/dataloader.py +++ b/deepmd/pd/utils/dataloader.py @@ -167,7 +167,7 @@ def construct_dataset(system): else: self.batch_sizes = batch_size * np.ones(len(systems), dtype=int) assert len(self.systems) == len(self.batch_sizes) - for system, batch_size in zip(self.systems, self.batch_sizes): + for system, batch_size in zip(self.systems, self.batch_sizes, strict=False): if dist.is_available() and dist.is_initialized(): system_batch_sampler = DistributedBatchSampler( system, diff --git a/deepmd/pd/utils/nlist.py b/deepmd/pd/utils/nlist.py index cd041a345c..203bdff46f 100644 --- a/deepmd/pd/utils/nlist.py +++ b/deepmd/pd/utils/nlist.py @@ -404,7 +404,7 @@ def build_multiple_neighbor_list( rr.masked_fill(nlist_mask, float("inf")) nlist0 = nlist ret = {} - for rc, ns in zip(rcuts[::-1], nsels[::-1]): + for rc, ns in zip(rcuts[::-1], nsels[::-1], strict=False): nlist0 = nlist0[:, :, :ns].masked_fill(rr[:, :, :ns] > rc, -1) ret[get_multiple_nlist_key(rc, ns)] = nlist0 return ret diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 2726b61152..03f6684c6a 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -376,7 +376,7 @@ def eval( return dict( zip( [x.name for x in request_defs], - out, + out, strict=False, ) ) diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 96b3baf6ec..4208bbe10c 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -274,12 +274,12 @@ def forward_atomic( raw_nlists = [ nlists[get_multiple_nlist_key(rcut, sel)] - for rcut, sel in zip(self.get_model_rcuts(), self.get_model_nsels()) + for rcut, sel in zip(self.get_model_rcuts(), self.get_model_nsels(), strict=False) ] nlists_ = [ nl if mt else nlist_distinguish_types(nl, extended_atype, sel) for mt, nl, sel in zip( - self.mixed_types_list, raw_nlists, self.get_model_sels() + self.mixed_types_list, raw_nlists, self.get_model_sels(), strict=False ) ] ener_list = [] diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index 0f001bc4c8..71d5dd277d 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -95,7 +95,7 @@ def __init__( start_idx = np.cumsum(np.pad(hybrid_sel, (1, 0), "constant"))[:-1] end_idx = start_idx + np.array(sub_sel) cut_idx = np.concatenate( - [range(ss, ee) for ss, ee in zip(start_idx, end_idx)] + [range(ss, ee) for ss, ee in zip(start_idx, end_idx, strict=False)] ).astype(np.int64) self.nlist_cut_idx.append(to_torch_tensor(cut_idx)) diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 59c165ddb0..67e451ac9c 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -785,7 +785,7 @@ def forward( # nfnl x nnei exclude_mask = self.emask(nlist, extended_atype).view(nfnl, self.nnei) for embedding_idx, (ll, compress_data_ii, compress_info_ii) in enumerate( - zip(self.filter_layers.networks, self.compress_data, self.compress_info) + zip(self.filter_layers.networks, self.compress_data, self.compress_info, strict=False) ): if self.type_one_side: ii = embedding_idx diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index 1d5a8fc1a8..acd632bbbc 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -496,7 +496,7 @@ def forward( exclude_mask = self.emask(nlist, atype_ext).view(nfnl, self.nnei) xyz_scatter_total = [] for ii, (ll, compress_data_ii, compress_info_ii) in enumerate( - zip(self.filter_layers.networks, self.compress_data, self.compress_info) + zip(self.filter_layers.networks, self.compress_data, self.compress_info, strict=False) ): # nfnl x nt mm = exclude_mask[:, self.sec[ii] : self.sec[ii + 1]] diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index 9771d8fe6f..76274a05c1 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -845,7 +845,7 @@ def forward( # nfnl x nnei exclude_mask = self.emask(nlist, extended_atype).view(nfnl, self.nnei) for embedding_idx, (ll, compress_data_ii, compress_info_ii) in enumerate( - zip(self.filter_layers.networks, self.compress_data, self.compress_info) + zip(self.filter_layers.networks, self.compress_data, self.compress_info, strict=False) ): ti = embedding_idx % self.ntypes nei_type_j = self.sel[ti] diff --git a/deepmd/pt/model/model/transform_output.py b/deepmd/pt/model/model/transform_output.py index 839a028fcf..467345c30f 100644 --- a/deepmd/pt/model/model/transform_output.py +++ b/deepmd/pt/model/model/transform_output.py @@ -122,7 +122,7 @@ def take_deriv( split_vv1 = torch.split(vv1, [1] * size, dim=-1) split_svv1 = torch.split(svv1, [1] * size, dim=-1) split_ff, split_avir = [], [] - for vvi, svvi in zip(split_vv1, split_svv1): + for vvi, svvi in zip(split_vv1, split_svv1, strict=False): # nf x nloc x 3, nf x nloc x 9 ffi, aviri = task_deriv_one( vvi, diff --git a/deepmd/pt/optimizer/LKF.py b/deepmd/pt/optimizer/LKF.py index c79e875f3e..9cad12d89d 100644 --- a/deepmd/pt/optimizer/LKF.py +++ b/deepmd/pt/optimizer/LKF.py @@ -292,7 +292,7 @@ def step(self, error: torch.Tensor) -> None: tmp = self.__split_weights(tmp) tmp_grad = self.__split_weights(tmp_grad) - for split_grad, split_weight in zip(tmp_grad, tmp): + for split_grad, split_weight in zip(tmp_grad, tmp, strict=False): nelement = split_grad.nelement() if param_sum == 0: diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 24440e19de..d940896c25 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -650,7 +650,7 @@ def single_model_finetune( self.wrapper.share_params( shared_links, resume=(resuming and not self.finetune_update_stat) or self.rank != 0, - model_key_prob_map=dict(zip(self.model_keys, self.model_prob)), + model_key_prob_map=dict(zip(self.model_keys, self.model_prob, strict=False)), data_stat_protect=_data_stat_protect[0], ) diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py index c991e59daa..ebab848564 100644 --- a/deepmd/pt/utils/dataloader.py +++ b/deepmd/pt/utils/dataloader.py @@ -157,7 +157,7 @@ def construct_dataset(system: str) -> DeepmdDataSetForLoader: else: self.batch_sizes = batch_size * np.ones(len(systems), dtype=int) assert len(self.systems) == len(self.batch_sizes) - for system, batch_size in zip(self.systems, self.batch_sizes): + for system, batch_size in zip(self.systems, self.batch_sizes, strict=False): if dist.is_available() and dist.is_initialized(): system_sampler = DistributedSampler(system) self.sampler_list.append(system_sampler) diff --git a/deepmd/pt/utils/nlist.py b/deepmd/pt/utils/nlist.py index 7f74e65f26..21b678ca1b 100644 --- a/deepmd/pt/utils/nlist.py +++ b/deepmd/pt/utils/nlist.py @@ -398,7 +398,7 @@ def build_multiple_neighbor_list( rr.masked_fill(nlist_mask, float("inf")) nlist0 = nlist ret = {} - for rc, ns in zip(rcuts[::-1], nsels[::-1]): + for rc, ns in zip(rcuts[::-1], nsels[::-1], strict=False): nlist0 = nlist0[:, :, :ns].masked_fill(rr[:, :, :ns] > rc, -1) ret[get_multiple_nlist_key(rc, ns)] = nlist0 return ret diff --git a/deepmd/tf/descriptor/loc_frame.py b/deepmd/tf/descriptor/loc_frame.py index 87db875038..1fd7f521e5 100644 --- a/deepmd/tf/descriptor/loc_frame.py +++ b/deepmd/tf/descriptor/loc_frame.py @@ -187,7 +187,7 @@ def compute_input_stats( sumn = [] sumv2 = [] for cc, bb, tt, nn, mm in zip( - data_coord, data_box, data_atype, natoms_vec, mesh + data_coord, data_box, data_atype, natoms_vec, mesh, strict=False ): sysv, sysv2, sysn = self._compute_dstats_sys_nonsmth(cc, bb, tt, nn, mm) sumv.append(sysv) diff --git a/deepmd/tf/descriptor/se_a.py b/deepmd/tf/descriptor/se_a.py index 3d2ab39c9f..52b38bbcad 100644 --- a/deepmd/tf/descriptor/se_a.py +++ b/deepmd/tf/descriptor/se_a.py @@ -374,7 +374,7 @@ def compute_input_stats( sumr2 = [] suma2 = [] for cc, bb, tt, nn, mm in zip( - data_coord, data_box, data_atype, natoms_vec, mesh + data_coord, data_box, data_atype, natoms_vec, mesh, strict=False ): sysr, sysr2, sysa, sysa2, sysn = self._compute_dstats_sys_smth( cc, bb, tt, nn, mm @@ -1331,7 +1331,7 @@ def init_variables( start_index_old[0] = 0 for nn, oo, ii, jj in zip( - n_descpt, n_descpt_old, start_index, start_index_old + n_descpt, n_descpt_old, start_index, start_index_old, strict=False ): if nn < oo: # new size is smaller, copy part of std diff --git a/deepmd/tf/descriptor/se_a_ef.py b/deepmd/tf/descriptor/se_a_ef.py index 37aa830431..835afdba5e 100644 --- a/deepmd/tf/descriptor/se_a_ef.py +++ b/deepmd/tf/descriptor/se_a_ef.py @@ -419,7 +419,7 @@ def compute_input_stats( sumr2 = [] suma2 = [] for cc, bb, tt, nn, mm, ee in zip( - data_coord, data_box, data_atype, natoms_vec, mesh, data_efield + data_coord, data_box, data_atype, natoms_vec, mesh, data_efield, strict=False ): sysr, sysr2, sysa, sysa2, sysn = self._compute_dstats_sys_smth( cc, bb, tt, nn, mm, ee diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index 3978709715..9f4277b670 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -379,7 +379,7 @@ def compute_input_stats( if mixed_type: sys_num = 0 for cc, bb, tt, nn, mm, r_n in zip( - data_coord, data_box, data_atype, natoms_vec, mesh, real_natoms_vec + data_coord, data_box, data_atype, natoms_vec, mesh, real_natoms_vec, strict=False ): sysr, sysr2, sysa, sysa2, sysn = self._compute_dstats_sys_smth( cc, bb, tt, nn, mm, mixed_type, r_n @@ -392,7 +392,7 @@ def compute_input_stats( suma2.append(sysa2) else: for cc, bb, tt, nn, mm in zip( - data_coord, data_box, data_atype, natoms_vec, mesh + data_coord, data_box, data_atype, natoms_vec, mesh, strict=False ): sysr, sysr2, sysa, sysa2, sysn = self._compute_dstats_sys_smth( cc, bb, tt, nn, mm diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py index c38a13d35a..9bf553326f 100644 --- a/deepmd/tf/descriptor/se_r.py +++ b/deepmd/tf/descriptor/se_r.py @@ -271,7 +271,7 @@ def compute_input_stats( sumn = [] sumr2 = [] for cc, bb, tt, nn, mm in zip( - data_coord, data_box, data_atype, natoms_vec, mesh + data_coord, data_box, data_atype, natoms_vec, mesh, strict=False ): sysr, sysr2, sysn = self._compute_dstats_sys_se_r(cc, bb, tt, nn, mm) sumr.append(sysr) diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index b03746a9c6..96fc5d9189 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -258,7 +258,7 @@ def compute_input_stats( sumr2 = [] suma2 = [] for cc, bb, tt, nn, mm in zip( - data_coord, data_box, data_atype, natoms_vec, mesh + data_coord, data_box, data_atype, natoms_vec, mesh, strict=False ): sysr, sysr2, sysa, sysa2, sysn = self._compute_dstats_sys_smth( cc, bb, tt, nn, mm diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index d1eca83055..00756f9c73 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -760,7 +760,7 @@ def eval( output = (output,) output_dict = { - odef.name: oo for oo, odef in zip(output, self.output_def.var_defs.values()) + odef.name: oo for oo, odef in zip(output, self.output_def.var_defs.values(), strict=False) } # ugly!! if self.modifier_type is not None and issubclass(self.model_type, DeepPot): diff --git a/deepmd/tf/nvnmd/entrypoints/mapt.py b/deepmd/tf/nvnmd/entrypoints/mapt.py index 121263974a..7e48bde300 100644 --- a/deepmd/tf/nvnmd/entrypoints/mapt.py +++ b/deepmd/tf/nvnmd/entrypoints/mapt.py @@ -439,11 +439,11 @@ def run_u2s(self): u = N2 * np.reshape(np.arange(0, N + 1) / N, [-1, 1]) # pylint: disable=no-explicit-dtype res_lst = run_sess(sess, vals, feed_dict={dic_ph["u"]: u}) - res_dic = dict(zip(keys, res_lst)) + res_dic = dict(zip(keys, res_lst, strict=False)) u2 = N2 * np.reshape(np.arange(0, N * 16 + 1) / (N * 16), [-1, 1]) # pylint: disable=no-explicit-dtype res_lst2 = run_sess(sess, vals, feed_dict={dic_ph["u"]: u2}) - res_dic2 = dict(zip(keys, res_lst2)) # reference for compare + res_dic2 = dict(zip(keys, res_lst2, strict=False)) # reference for compare # change value for tt in range(ndim): @@ -536,11 +536,11 @@ def run_s2g(self): s = N2 * np.reshape(np.arange(0, N + 1) / N, [-1, 1]) + smin_ # pylint: disable=no-explicit-dtype res_lst = run_sess(sess, vals, feed_dict={dic_ph["s"]: s}) - res_dic = dict(zip(keys, res_lst)) + res_dic = dict(zip(keys, res_lst, strict=False)) s2 = N2 * np.reshape(np.arange(0, N * 16 + 1) / (N * 16), [-1, 1]) + smin_ # pylint: disable=no-explicit-dtype res_lst2 = run_sess(sess, vals, feed_dict={dic_ph["s"]: s2}) - res_dic2 = dict(zip(keys, res_lst2)) + res_dic2 = dict(zip(keys, res_lst2, strict=False)) sess.close() return res_dic, res_dic2 @@ -601,7 +601,7 @@ def run_t2g(self): vals = list(dic_ph.values()) # res_lst = run_sess(sess, vals, feed_dict={}) - res_dic = dict(zip(keys, res_lst)) + res_dic = dict(zip(keys, res_lst, strict=False)) sess.close() return res_dic diff --git a/deepmd/tf/train/trainer.py b/deepmd/tf/train/trainer.py index f70c919301..e10cff967a 100644 --- a/deepmd/tf/train/trainer.py +++ b/deepmd/tf/train/trainer.py @@ -967,4 +967,4 @@ def get_data_dict(self, batch_list: list[np.ndarray]) -> dict[str, np.ndarray]: dict[str, np.ndarray] The dict of the loaded data. """ - return dict(zip(self.data_keys, batch_list)) + return dict(zip(self.data_keys, batch_list, strict=False)) diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py index 012ffb4260..effd3bd53f 100644 --- a/deepmd/utils/batch_size.py +++ b/deepmd/utils/batch_size.py @@ -248,7 +248,7 @@ def concate_result(r: list[Any]) -> Any: return ret if not returned_dict: - r_list = [concate_result(r) for r in zip(*results)] + r_list = [concate_result(r) for r in zip(*results, strict=False)] r = tuple(r_list) if len(r) == 1: # avoid returning tuple if callable doesn't return tuple diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 82ef5ec25d..c59230c42f 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -172,7 +172,7 @@ def __init__( rule = int(words[1]) filtered_data_systems = [] filtered_system_dirs = [] - for sys_dir, data_sys in zip(self.system_dirs, self.data_systems): + for sys_dir, data_sys in zip(self.system_dirs, self.data_systems, strict=False): if data_sys.get_natoms() <= rule: filtered_data_systems.append(data_sys) filtered_system_dirs.append(sys_dir) diff --git a/deepmd/utils/model_branch_dict.py b/deepmd/utils/model_branch_dict.py index 501f5287b6..3c2c37a134 100644 --- a/deepmd/utils/model_branch_dict.py +++ b/deepmd/utils/model_branch_dict.py @@ -169,7 +169,7 @@ def as_table(self) -> str: # Step 3: Determine actual width for each column # For the first two columns, we already decided the exact widths above. col_widths: list[int] = [] - for idx, col in enumerate(zip(*wrapped_rows)): + for idx, col in enumerate(zip(*wrapped_rows, strict=False)): if idx == 0: col_widths.append(branch_col_width) elif idx == 1: @@ -187,7 +187,7 @@ def draw_row_line(cells_parts: list[list[str]]) -> str: return ( "| " + " | ".join( - part.ljust(width) for part, width in zip(cells_parts, col_widths) + part.ljust(width) for part, width in zip(cells_parts, col_widths, strict=False) ) + " |" ) diff --git a/deepmd/utils/update_sel.py b/deepmd/utils/update_sel.py index 616b69c2ca..75e62f5006 100644 --- a/deepmd/utils/update_sel.py +++ b/deepmd/utils/update_sel.py @@ -43,7 +43,7 @@ def update_one_sel( sel = [int(self.wrap_up_4(ii * ratio)) for ii in tmp_sel] else: # sel is set by user - for ii, (tt, dd) in enumerate(zip(tmp_sel, sel)): + for ii, (tt, dd) in enumerate(zip(tmp_sel, sel, strict=False)): if dd and tt > dd: # we may skip warning for sel=0, where the user is likely # to exclude such type in the descriptor diff --git a/source/tests/common/dpmodel/test_pairtab_atomic_model.py b/source/tests/common/dpmodel/test_pairtab_atomic_model.py index f8145c258a..7e9affbf09 100644 --- a/source/tests/common/dpmodel/test_pairtab_atomic_model.py +++ b/source/tests/common/dpmodel/test_pairtab_atomic_model.py @@ -157,7 +157,7 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: 0.03, 0.035, 0.025, - ], + ], strict=False, ): extended_coord = np.array( [ diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index c1e6644be6..42814566f3 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -355,7 +355,7 @@ def test_tf_consistent_with_ref(self) -> None: data2.pop("@version") np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): np.testing.assert_allclose( rr1.ravel(), rr2.ravel(), rtol=self.rtol, atol=self.atol ) @@ -372,7 +372,7 @@ def test_tf_self_consistent(self) -> None: obj2 = self.tf_class.deserialize(data1, suffix=self.unique_id) ret2, data2 = self.get_tf_ret_serialization_from_cls(obj2) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -391,7 +391,7 @@ def test_dp_consistent_with_ref(self) -> None: ret2 = self.extract_ret(ret2, self.RefBackend.DP) data2 = dp_obj.serialize() np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): if rr1 is SKIP_FLAG or rr2 is SKIP_FLAG: continue np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) @@ -407,7 +407,7 @@ def test_dp_self_consistent(self) -> None: obj1 = self.dp_class.deserialize(data1) ret2, data2 = self.get_dp_ret_serialization_from_cls(obj1) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): if isinstance(rr1, np.ndarray) and isinstance(rr2, np.ndarray): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -436,7 +436,7 @@ def test_pt_consistent_with_ref(self) -> None: data1.pop("@variables", None) data2.pop("@variables", None) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -449,7 +449,7 @@ def test_pt_self_consistent(self) -> None: obj2 = self.pt_class.deserialize(data1) ret2, data2 = self.get_pt_ret_serialization_from_cls(obj2) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): if isinstance(rr1, np.ndarray) and isinstance(rr2, np.ndarray): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -473,7 +473,7 @@ def test_jax_consistent_with_ref(self) -> None: data1.pop("@variables", None) data2.pop("@variables", None) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -486,7 +486,7 @@ def test_jax_self_consistent(self) -> None: obj1 = self.jax_class.deserialize(data1) ret2, data2 = self.get_jax_ret_serialization_from_cls(obj1) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): if isinstance(rr1, np.ndarray) and isinstance(rr2, np.ndarray): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -515,7 +515,7 @@ def test_pd_consistent_with_ref(self): data1.pop("@variables", None) data2.pop("@variables", None) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -528,7 +528,7 @@ def test_pd_self_consistent(self): obj2 = self.pd_class.deserialize(data1) ret2, data2 = self.get_pd_ret_serialization_from_cls(obj2) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): if isinstance(rr1, np.ndarray) and isinstance(rr2, np.ndarray): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -550,7 +550,7 @@ def test_array_api_strict_consistent_with_ref(self) -> None: ret2 = self.extract_ret(ret2, self.RefBackend.ARRAY_API_STRICT) data2 = array_api_strict_obj.serialize() np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -564,7 +564,7 @@ def test_array_api_strict_self_consistent(self) -> None: obj1 = self.array_api_strict_class.deserialize(data1) ret2, data2 = self.get_array_api_strict_ret_serialization_from_cls(obj1) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): if isinstance(rr1, np.ndarray) and isinstance(rr2, np.ndarray): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" diff --git a/source/tests/consistent/fitting/test_dipole.py b/source/tests/consistent/fitting/test_dipole.py index 010944d109..bf2907df03 100644 --- a/source/tests/consistent/fitting/test_dipole.py +++ b/source/tests/consistent/fitting/test_dipole.py @@ -285,7 +285,7 @@ def test_tf_consistent_with_ref(self) -> None: if networks is None: data1["nets"]["networks"][ii] = None np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2): + for rr1, rr2 in zip(ret1, ret2, strict=False): np.testing.assert_allclose( rr1.ravel()[: rr2.size], rr2.ravel(), rtol=self.rtol, atol=self.atol ) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 7ae9af6891..cc7b3dcde1 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -202,14 +202,14 @@ def test_deep_eval(self) -> None: ) rets_nopbc.append(ret) for ret in rets[1:]: - for vv1, vv2 in zip(rets[0], ret): + for vv1, vv2 in zip(rets[0], ret, strict=False): if np.isnan(vv2).all(): # expect all nan if not supported continue np.testing.assert_allclose(vv1, vv2, rtol=1e-12, atol=1e-12) for idx, ret in enumerate(rets_nopbc[1:]): - for vv1, vv2 in zip(rets_nopbc[0], ret): + for vv1, vv2 in zip(rets_nopbc[0], ret, strict=False): if np.isnan(vv2).all(): # expect all nan if not supported continue diff --git a/source/tests/pd/model/test_nlist.py b/source/tests/pd/model/test_nlist.py index 0947355ac0..fc73e7ee8f 100644 --- a/source/tests/pd/model/test_nlist.py +++ b/source/tests/pd/model/test_nlist.py @@ -253,7 +253,7 @@ def test_build_directional_nlist(self): ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) - for distinguish_types, mysel in zip([True, False], [sum(self.nsel), 300]): + for distinguish_types, mysel in zip([True, False], [sum(self.nsel), 300], strict=False): # full neighbor list nlist_full = build_neighbor_list( ecoord, diff --git a/source/tests/pd/model/test_se_e2_a.py b/source/tests/pd/model/test_se_e2_a.py index b1e6abe5ae..c7809586e6 100644 --- a/source/tests/pd/model/test_se_e2_a.py +++ b/source/tests/pd/model/test_se_e2_a.py @@ -96,7 +96,7 @@ def test_consistency( self.atype_ext, self.nlist, ) - for aa, bb in zip([rd1, gr1, sw1], [rd2, gr2, sw2]): + for aa, bb in zip([rd1, gr1, sw1], [rd2, gr2, sw2], strict=False): np.testing.assert_allclose( aa.detach().cpu().numpy(), bb, diff --git a/source/tests/pd/test_utils.py b/source/tests/pd/test_utils.py index 8d25cff964..bfbff0dd8f 100644 --- a/source/tests/pd/test_utils.py +++ b/source/tests/pd/test_utils.py @@ -20,7 +20,7 @@ def test_to_numpy(self): foo = rng.normal([3, 4]) for ptp, npp in zip( [paddle.float16, paddle.float32, paddle.float64], - [np.float16, np.float32, np.float64], + [np.float16, np.float32, np.float64], strict=False, ): foo = foo.astype(npp) bar = to_paddle_tensor(foo) diff --git a/source/tests/pt/model/test_descriptor_se_r.py b/source/tests/pt/model/test_descriptor_se_r.py index 8c6e2c6776..86722bc19b 100644 --- a/source/tests/pt/model/test_descriptor_se_r.py +++ b/source/tests/pt/model/test_descriptor_se_r.py @@ -100,7 +100,7 @@ def test_consistency( self.atype_ext, self.nlist, ) - for aa, bb in zip([rd1, sw1], [rd2, sw2]): + for aa, bb in zip([rd1, sw1], [rd2, sw2], strict=False): np.testing.assert_allclose( aa.detach().cpu().numpy(), bb, diff --git a/source/tests/pt/model/test_nlist.py b/source/tests/pt/model/test_nlist.py index fdcbb813bc..24aee9d4b9 100644 --- a/source/tests/pt/model/test_nlist.py +++ b/source/tests/pt/model/test_nlist.py @@ -245,7 +245,7 @@ def test_build_directional_nlist(self) -> None: ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) - for distinguish_types, mysel in zip([True, False], [sum(self.nsel), 300]): + for distinguish_types, mysel in zip([True, False], [sum(self.nsel), 300], strict=False): # full neighbor list nlist_full = build_neighbor_list( ecoord, diff --git a/source/tests/pt/model/test_pairtab_atomic_model.py b/source/tests/pt/model/test_pairtab_atomic_model.py index 0a6f6ea933..f41e47a9d7 100644 --- a/source/tests/pt/model/test_pairtab_atomic_model.py +++ b/source/tests/pt/model/test_pairtab_atomic_model.py @@ -218,7 +218,7 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: 0.03, 0.035, 0.025, - ], + ], strict=False, ): extended_coord = torch.tensor( [ diff --git a/source/tests/pt/model/test_se_e2_a.py b/source/tests/pt/model/test_se_e2_a.py index bdedd440cb..d87f5c8348 100644 --- a/source/tests/pt/model/test_se_e2_a.py +++ b/source/tests/pt/model/test_se_e2_a.py @@ -96,7 +96,7 @@ def test_consistency( self.atype_ext, self.nlist, ) - for aa, bb in zip([rd1, gr1, sw1], [rd2, gr2, sw2]): + for aa, bb in zip([rd1, gr1, sw1], [rd2, gr2, sw2], strict=False): np.testing.assert_allclose( aa.detach().cpu().numpy(), bb, diff --git a/source/tests/pt/model/test_se_t.py b/source/tests/pt/model/test_se_t.py index fb60ca011e..cddb455295 100644 --- a/source/tests/pt/model/test_se_t.py +++ b/source/tests/pt/model/test_se_t.py @@ -100,7 +100,7 @@ def test_consistency( self.atype_ext, self.nlist, ) - for aa, bb in zip([rd1, sw1], [rd2, sw2]): + for aa, bb in zip([rd1, sw1], [rd2, sw2], strict=False): np.testing.assert_allclose( aa.detach().cpu().numpy(), bb, diff --git a/source/tests/pt/test_utils.py b/source/tests/pt/test_utils.py index 42bf2679a8..bbfa4a4294 100644 --- a/source/tests/pt/test_utils.py +++ b/source/tests/pt/test_utils.py @@ -20,7 +20,7 @@ def test_to_numpy(self) -> None: foo = rng.normal([3, 4]) for ptp, npp in zip( [torch.float16, torch.float32, torch.float64], - [np.float16, np.float32, np.float64], + [np.float16, np.float32, np.float64], strict=False, ): foo = foo.astype(npp) bar = to_torch_tensor(foo) diff --git a/source/tests/tf/test_linear_model.py b/source/tests/tf/test_linear_model.py index e99c1042dd..5746eab84c 100644 --- a/source/tests/tf/test_linear_model.py +++ b/source/tests/tf/test_linear_model.py @@ -37,7 +37,7 @@ def setUp(self) -> None: os.path.join(infer_path, "deeppot-1.pbtxt"), ] self.graph_dirs = [pbtxt.replace("pbtxt", "pb") for pbtxt in self.pbtxts] - for pbtxt, pb in zip(self.pbtxts, self.graph_dirs): + for pbtxt, pb in zip(self.pbtxts, self.graph_dirs, strict=False): convert_pbtxt_to_pb(pbtxt, pb) self.graphs = [DeepPotential(pb) for pb in self.graph_dirs] diff --git a/source/tests/tf/test_model_devi.py b/source/tests/tf/test_model_devi.py index be9486cf02..de02dedba5 100644 --- a/source/tests/tf/test_model_devi.py +++ b/source/tests/tf/test_model_devi.py @@ -41,7 +41,7 @@ def setUp(self) -> None: os.path.join(infer_path, "deeppot-1.pbtxt"), ] self.graph_dirs = [pbtxt.replace("pbtxt", "pb") for pbtxt in self.pbtxts] - for pbtxt, pb in zip(self.pbtxts, self.graph_dirs): + for pbtxt, pb in zip(self.pbtxts, self.graph_dirs, strict=False): convert_pbtxt_to_pb(pbtxt, pb) self.graphs = [DeepPotential(pb) for pb in self.graph_dirs] self.output = os.path.join(tests_path, "model_devi.out") @@ -216,7 +216,7 @@ def setUpClass(cls) -> None: os.path.join(infer_path, "fparam_aparam.pbtxt"), ] cls.graph_dirs = [pbtxt.replace("pbtxt", "pb") for pbtxt in cls.pbtxts] - for pbtxt, pb in zip(cls.pbtxts, cls.graph_dirs): + for pbtxt, pb in zip(cls.pbtxts, cls.graph_dirs, strict=False): convert_pbtxt_to_pb(pbtxt, pb) cls.graphs = [DeepPotential(pb) for pb in cls.graph_dirs] diff --git a/source/tests/tf/test_model_devi_mix.py b/source/tests/tf/test_model_devi_mix.py index 4ef9f39cab..b3907751aa 100644 --- a/source/tests/tf/test_model_devi_mix.py +++ b/source/tests/tf/test_model_devi_mix.py @@ -58,7 +58,7 @@ def setUp(self) -> None: os.path.join(infer_path, "se_atten_no_atten_2.pbtxt"), ] self.graph_dirs = [pbtxt.replace("pbtxt", "pb") for pbtxt in self.pbtxts] - for pbtxt, pb in zip(self.pbtxts, self.graph_dirs): + for pbtxt, pb in zip(self.pbtxts, self.graph_dirs, strict=False): convert_pbtxt_to_pb(pbtxt, pb) self.graphs = [DeepPotential(pb) for pb in self.graph_dirs] self.output = os.path.join(tests_path, "model_devi.out") diff --git a/source/tests/tf/test_nvnmd_entrypoints.py b/source/tests/tf/test_nvnmd_entrypoints.py index 6833b4e0f8..2e3c3f2f37 100644 --- a/source/tests/tf/test_nvnmd_entrypoints.py +++ b/source/tests/tf/test_nvnmd_entrypoints.py @@ -466,7 +466,7 @@ def test_model_qnn_v0(self) -> None: tensordic[key] = graph.get_tensor_by_name(key + ":0") # get value valuelist = sess.run(list(tensordic.values()), feed_dict=feed_dict) - valuedic = dict(zip(tensordic.keys(), valuelist)) + valuedic = dict(zip(tensordic.keys(), valuelist, strict=False)) # test # o_descriptor idx = np.array([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]) @@ -783,7 +783,7 @@ def test_model_qnn_v1(self) -> None: tensordic[key] = graph.get_tensor_by_name(key + ":0") # get value valuelist = sess.run(list(tensordic.values()), feed_dict=feed_dict) - valuedic = dict(zip(tensordic.keys(), valuelist)) + valuedic = dict(zip(tensordic.keys(), valuelist, strict=False)) # test # o_descriptor idx = np.array([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]) From f8aae868ac1ad52796c6b8a4207ed160b264ae7e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 7 Jan 2026 17:56:55 +0000 Subject: [PATCH 03/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/atomic_model/linear_atomic_model.py | 4 +++- deepmd/dpmodel/infer/deep_eval.py | 3 ++- deepmd/jax/infer/deep_eval.py | 3 ++- deepmd/pd/infer/deep_eval.py | 3 ++- deepmd/pd/model/descriptor/se_a.py | 7 ++++++- deepmd/pt/infer/deep_eval.py | 3 ++- deepmd/pt/model/atomic_model/linear_atomic_model.py | 4 +++- deepmd/pt/model/descriptor/se_a.py | 7 ++++++- deepmd/pt/model/descriptor/se_r.py | 7 ++++++- deepmd/pt/model/descriptor/se_t.py | 7 ++++++- deepmd/pt/train/training.py | 4 +++- deepmd/tf/descriptor/se_a_ef.py | 8 +++++++- deepmd/tf/descriptor/se_atten.py | 8 +++++++- deepmd/tf/infer/deep_eval.py | 3 ++- deepmd/utils/data_system.py | 4 +++- deepmd/utils/model_branch_dict.py | 3 ++- source/tests/common/dpmodel/test_pairtab_atomic_model.py | 3 ++- source/tests/pd/model/test_nlist.py | 4 +++- source/tests/pd/test_utils.py | 3 ++- source/tests/pt/model/test_nlist.py | 4 +++- source/tests/pt/model/test_pairtab_atomic_model.py | 3 ++- source/tests/pt/test_utils.py | 3 ++- 22 files changed, 76 insertions(+), 22 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 743ca1cf39..766f8a1821 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -235,7 +235,9 @@ def forward_atomic( ) raw_nlists = [ nlists[get_multiple_nlist_key(rcut, sel)] - for rcut, sel in zip(self.get_model_rcuts(), self.get_model_nsels(), strict=False) + for rcut, sel in zip( + self.get_model_rcuts(), self.get_model_nsels(), strict=False + ) ] nlists_ = [ nl if mt else nlist_distinguish_types(nl, extended_atype, sel) diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index 4526680fc4..bb366a84ca 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -228,7 +228,8 @@ def eval( return dict( zip( [x.name for x in request_defs], - out, strict=False, + out, + strict=False, ) ) diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 1ae9662080..8a2d5efef6 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -248,7 +248,8 @@ def eval( return dict( zip( [x.name for x in request_defs], - out, strict=False, + out, + strict=False, ) ) diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index d40f91cb8d..fcf135741b 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -384,7 +384,8 @@ def eval( return dict( zip( [x.name for x in request_defs], - out, strict=False, + out, + strict=False, ) ) diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py index dbed543041..7d626425ec 100644 --- a/deepmd/pd/model/descriptor/se_a.py +++ b/deepmd/pd/model/descriptor/se_a.py @@ -763,7 +763,12 @@ def forward( # nfnl x nnei exclude_mask = self.emask(nlist, extended_atype).reshape([nfnl, self.nnei]) for embedding_idx, (ll, compress_data_ii, compress_info_ii) in enumerate( - zip(self.filter_layers.networks, self.compress_data, self.compress_info, strict=False) + zip( + self.filter_layers.networks, + self.compress_data, + self.compress_info, + strict=False, + ) ): if self.type_one_side: ii = embedding_idx diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 03f6684c6a..dfe5f02895 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -376,7 +376,8 @@ def eval( return dict( zip( [x.name for x in request_defs], - out, strict=False, + out, + strict=False, ) ) diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 4208bbe10c..e6c7f5fa7f 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -274,7 +274,9 @@ def forward_atomic( raw_nlists = [ nlists[get_multiple_nlist_key(rcut, sel)] - for rcut, sel in zip(self.get_model_rcuts(), self.get_model_nsels(), strict=False) + for rcut, sel in zip( + self.get_model_rcuts(), self.get_model_nsels(), strict=False + ) ] nlists_ = [ nl if mt else nlist_distinguish_types(nl, extended_atype, sel) diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 67e451ac9c..40665bf5f2 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -785,7 +785,12 @@ def forward( # nfnl x nnei exclude_mask = self.emask(nlist, extended_atype).view(nfnl, self.nnei) for embedding_idx, (ll, compress_data_ii, compress_info_ii) in enumerate( - zip(self.filter_layers.networks, self.compress_data, self.compress_info, strict=False) + zip( + self.filter_layers.networks, + self.compress_data, + self.compress_info, + strict=False, + ) ): if self.type_one_side: ii = embedding_idx diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index acd632bbbc..fdb050968f 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -496,7 +496,12 @@ def forward( exclude_mask = self.emask(nlist, atype_ext).view(nfnl, self.nnei) xyz_scatter_total = [] for ii, (ll, compress_data_ii, compress_info_ii) in enumerate( - zip(self.filter_layers.networks, self.compress_data, self.compress_info, strict=False) + zip( + self.filter_layers.networks, + self.compress_data, + self.compress_info, + strict=False, + ) ): # nfnl x nt mm = exclude_mask[:, self.sec[ii] : self.sec[ii + 1]] diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index 76274a05c1..2db7ee18c7 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -845,7 +845,12 @@ def forward( # nfnl x nnei exclude_mask = self.emask(nlist, extended_atype).view(nfnl, self.nnei) for embedding_idx, (ll, compress_data_ii, compress_info_ii) in enumerate( - zip(self.filter_layers.networks, self.compress_data, self.compress_info, strict=False) + zip( + self.filter_layers.networks, + self.compress_data, + self.compress_info, + strict=False, + ) ): ti = embedding_idx % self.ntypes nei_type_j = self.sel[ti] diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index d940896c25..802b27726e 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -650,7 +650,9 @@ def single_model_finetune( self.wrapper.share_params( shared_links, resume=(resuming and not self.finetune_update_stat) or self.rank != 0, - model_key_prob_map=dict(zip(self.model_keys, self.model_prob, strict=False)), + model_key_prob_map=dict( + zip(self.model_keys, self.model_prob, strict=False) + ), data_stat_protect=_data_stat_protect[0], ) diff --git a/deepmd/tf/descriptor/se_a_ef.py b/deepmd/tf/descriptor/se_a_ef.py index 835afdba5e..146ffc424d 100644 --- a/deepmd/tf/descriptor/se_a_ef.py +++ b/deepmd/tf/descriptor/se_a_ef.py @@ -419,7 +419,13 @@ def compute_input_stats( sumr2 = [] suma2 = [] for cc, bb, tt, nn, mm, ee in zip( - data_coord, data_box, data_atype, natoms_vec, mesh, data_efield, strict=False + data_coord, + data_box, + data_atype, + natoms_vec, + mesh, + data_efield, + strict=False, ): sysr, sysr2, sysa, sysa2, sysn = self._compute_dstats_sys_smth( cc, bb, tt, nn, mm, ee diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index 9f4277b670..8cce8f7da7 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -379,7 +379,13 @@ def compute_input_stats( if mixed_type: sys_num = 0 for cc, bb, tt, nn, mm, r_n in zip( - data_coord, data_box, data_atype, natoms_vec, mesh, real_natoms_vec, strict=False + data_coord, + data_box, + data_atype, + natoms_vec, + mesh, + real_natoms_vec, + strict=False, ): sysr, sysr2, sysa, sysa2, sysn = self._compute_dstats_sys_smth( cc, bb, tt, nn, mm, mixed_type, r_n diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 00756f9c73..a7eedb2e16 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -760,7 +760,8 @@ def eval( output = (output,) output_dict = { - odef.name: oo for oo, odef in zip(output, self.output_def.var_defs.values(), strict=False) + odef.name: oo + for oo, odef in zip(output, self.output_def.var_defs.values(), strict=False) } # ugly!! if self.modifier_type is not None and issubclass(self.model_type, DeepPot): diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index c59230c42f..6c4fe4a190 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -172,7 +172,9 @@ def __init__( rule = int(words[1]) filtered_data_systems = [] filtered_system_dirs = [] - for sys_dir, data_sys in zip(self.system_dirs, self.data_systems, strict=False): + for sys_dir, data_sys in zip( + self.system_dirs, self.data_systems, strict=False + ): if data_sys.get_natoms() <= rule: filtered_data_systems.append(data_sys) filtered_system_dirs.append(sys_dir) diff --git a/deepmd/utils/model_branch_dict.py b/deepmd/utils/model_branch_dict.py index 3c2c37a134..01348c3cee 100644 --- a/deepmd/utils/model_branch_dict.py +++ b/deepmd/utils/model_branch_dict.py @@ -187,7 +187,8 @@ def draw_row_line(cells_parts: list[list[str]]) -> str: return ( "| " + " | ".join( - part.ljust(width) for part, width in zip(cells_parts, col_widths, strict=False) + part.ljust(width) + for part, width in zip(cells_parts, col_widths, strict=False) ) + " |" ) diff --git a/source/tests/common/dpmodel/test_pairtab_atomic_model.py b/source/tests/common/dpmodel/test_pairtab_atomic_model.py index 7e9affbf09..19cf21d13c 100644 --- a/source/tests/common/dpmodel/test_pairtab_atomic_model.py +++ b/source/tests/common/dpmodel/test_pairtab_atomic_model.py @@ -157,7 +157,8 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: 0.03, 0.035, 0.025, - ], strict=False, + ], + strict=False, ): extended_coord = np.array( [ diff --git a/source/tests/pd/model/test_nlist.py b/source/tests/pd/model/test_nlist.py index fc73e7ee8f..b45767aa52 100644 --- a/source/tests/pd/model/test_nlist.py +++ b/source/tests/pd/model/test_nlist.py @@ -253,7 +253,9 @@ def test_build_directional_nlist(self): ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) - for distinguish_types, mysel in zip([True, False], [sum(self.nsel), 300], strict=False): + for distinguish_types, mysel in zip( + [True, False], [sum(self.nsel), 300], strict=False + ): # full neighbor list nlist_full = build_neighbor_list( ecoord, diff --git a/source/tests/pd/test_utils.py b/source/tests/pd/test_utils.py index bfbff0dd8f..f4aafe97af 100644 --- a/source/tests/pd/test_utils.py +++ b/source/tests/pd/test_utils.py @@ -20,7 +20,8 @@ def test_to_numpy(self): foo = rng.normal([3, 4]) for ptp, npp in zip( [paddle.float16, paddle.float32, paddle.float64], - [np.float16, np.float32, np.float64], strict=False, + [np.float16, np.float32, np.float64], + strict=False, ): foo = foo.astype(npp) bar = to_paddle_tensor(foo) diff --git a/source/tests/pt/model/test_nlist.py b/source/tests/pt/model/test_nlist.py index 24aee9d4b9..cfe63d0fa7 100644 --- a/source/tests/pt/model/test_nlist.py +++ b/source/tests/pt/model/test_nlist.py @@ -245,7 +245,9 @@ def test_build_directional_nlist(self) -> None: ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) - for distinguish_types, mysel in zip([True, False], [sum(self.nsel), 300], strict=False): + for distinguish_types, mysel in zip( + [True, False], [sum(self.nsel), 300], strict=False + ): # full neighbor list nlist_full = build_neighbor_list( ecoord, diff --git a/source/tests/pt/model/test_pairtab_atomic_model.py b/source/tests/pt/model/test_pairtab_atomic_model.py index f41e47a9d7..bd9a629a10 100644 --- a/source/tests/pt/model/test_pairtab_atomic_model.py +++ b/source/tests/pt/model/test_pairtab_atomic_model.py @@ -218,7 +218,8 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: 0.03, 0.035, 0.025, - ], strict=False, + ], + strict=False, ): extended_coord = torch.tensor( [ diff --git a/source/tests/pt/test_utils.py b/source/tests/pt/test_utils.py index bbfa4a4294..33fe58b19b 100644 --- a/source/tests/pt/test_utils.py +++ b/source/tests/pt/test_utils.py @@ -20,7 +20,8 @@ def test_to_numpy(self) -> None: foo = rng.normal([3, 4]) for ptp, npp in zip( [torch.float16, torch.float32, torch.float64], - [np.float16, np.float32, np.float64], strict=False, + [np.float16, np.float32, np.float64], + strict=False, ): foo = foo.astype(npp) bar = to_torch_tensor(foo) From 9cb8913c9dc629aa907982899c7e9ce831aeb336 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 8 Jan 2026 03:19:22 +0800 Subject: [PATCH 04/12] exclude B905 in pytorch --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 1e156f77c0..e80089e6b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -430,7 +430,7 @@ runtime-evaluated-base-classes = ["torch.nn.Module"] "backend/**" = ["ANN"] "data/**" = ["ANN"] "deepmd/tf/**" = ["TID253", "ANN"] -"deepmd/pt/**" = ["TID253"] +"deepmd/pt/**" = ["TID253", "B905"] "deepmd/jax/**" = ["TID253"] "deepmd/pd/**" = ["TID253", "ANN"] From 0c373c73cebbbd9d138a6d69b901cbe44a7683a6 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 8 Jan 2026 03:19:37 +0800 Subject: [PATCH 05/12] revert pytorch for torchscript --- deepmd/pt/infer/deep_eval.py | 2 +- deepmd/pt/model/atomic_model/linear_atomic_model.py | 4 ++-- deepmd/pt/model/descriptor/hybrid.py | 2 +- deepmd/pt/model/descriptor/se_a.py | 2 +- deepmd/pt/model/descriptor/se_r.py | 2 +- deepmd/pt/model/descriptor/se_t.py | 2 +- deepmd/pt/model/model/transform_output.py | 2 +- deepmd/pt/optimizer/LKF.py | 2 +- deepmd/pt/train/training.py | 2 +- deepmd/pt/utils/dataloader.py | 2 +- deepmd/pt/utils/nlist.py | 2 +- 11 files changed, 12 insertions(+), 12 deletions(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 03f6684c6a..2726b61152 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -376,7 +376,7 @@ def eval( return dict( zip( [x.name for x in request_defs], - out, strict=False, + out, ) ) diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 4208bbe10c..96b3baf6ec 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -274,12 +274,12 @@ def forward_atomic( raw_nlists = [ nlists[get_multiple_nlist_key(rcut, sel)] - for rcut, sel in zip(self.get_model_rcuts(), self.get_model_nsels(), strict=False) + for rcut, sel in zip(self.get_model_rcuts(), self.get_model_nsels()) ] nlists_ = [ nl if mt else nlist_distinguish_types(nl, extended_atype, sel) for mt, nl, sel in zip( - self.mixed_types_list, raw_nlists, self.get_model_sels(), strict=False + self.mixed_types_list, raw_nlists, self.get_model_sels() ) ] ener_list = [] diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index 71d5dd277d..0f001bc4c8 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -95,7 +95,7 @@ def __init__( start_idx = np.cumsum(np.pad(hybrid_sel, (1, 0), "constant"))[:-1] end_idx = start_idx + np.array(sub_sel) cut_idx = np.concatenate( - [range(ss, ee) for ss, ee in zip(start_idx, end_idx, strict=False)] + [range(ss, ee) for ss, ee in zip(start_idx, end_idx)] ).astype(np.int64) self.nlist_cut_idx.append(to_torch_tensor(cut_idx)) diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 67e451ac9c..59c165ddb0 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -785,7 +785,7 @@ def forward( # nfnl x nnei exclude_mask = self.emask(nlist, extended_atype).view(nfnl, self.nnei) for embedding_idx, (ll, compress_data_ii, compress_info_ii) in enumerate( - zip(self.filter_layers.networks, self.compress_data, self.compress_info, strict=False) + zip(self.filter_layers.networks, self.compress_data, self.compress_info) ): if self.type_one_side: ii = embedding_idx diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index acd632bbbc..1d5a8fc1a8 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -496,7 +496,7 @@ def forward( exclude_mask = self.emask(nlist, atype_ext).view(nfnl, self.nnei) xyz_scatter_total = [] for ii, (ll, compress_data_ii, compress_info_ii) in enumerate( - zip(self.filter_layers.networks, self.compress_data, self.compress_info, strict=False) + zip(self.filter_layers.networks, self.compress_data, self.compress_info) ): # nfnl x nt mm = exclude_mask[:, self.sec[ii] : self.sec[ii + 1]] diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index 76274a05c1..9771d8fe6f 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -845,7 +845,7 @@ def forward( # nfnl x nnei exclude_mask = self.emask(nlist, extended_atype).view(nfnl, self.nnei) for embedding_idx, (ll, compress_data_ii, compress_info_ii) in enumerate( - zip(self.filter_layers.networks, self.compress_data, self.compress_info, strict=False) + zip(self.filter_layers.networks, self.compress_data, self.compress_info) ): ti = embedding_idx % self.ntypes nei_type_j = self.sel[ti] diff --git a/deepmd/pt/model/model/transform_output.py b/deepmd/pt/model/model/transform_output.py index 467345c30f..839a028fcf 100644 --- a/deepmd/pt/model/model/transform_output.py +++ b/deepmd/pt/model/model/transform_output.py @@ -122,7 +122,7 @@ def take_deriv( split_vv1 = torch.split(vv1, [1] * size, dim=-1) split_svv1 = torch.split(svv1, [1] * size, dim=-1) split_ff, split_avir = [], [] - for vvi, svvi in zip(split_vv1, split_svv1, strict=False): + for vvi, svvi in zip(split_vv1, split_svv1): # nf x nloc x 3, nf x nloc x 9 ffi, aviri = task_deriv_one( vvi, diff --git a/deepmd/pt/optimizer/LKF.py b/deepmd/pt/optimizer/LKF.py index 9cad12d89d..c79e875f3e 100644 --- a/deepmd/pt/optimizer/LKF.py +++ b/deepmd/pt/optimizer/LKF.py @@ -292,7 +292,7 @@ def step(self, error: torch.Tensor) -> None: tmp = self.__split_weights(tmp) tmp_grad = self.__split_weights(tmp_grad) - for split_grad, split_weight in zip(tmp_grad, tmp, strict=False): + for split_grad, split_weight in zip(tmp_grad, tmp): nelement = split_grad.nelement() if param_sum == 0: diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index d940896c25..24440e19de 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -650,7 +650,7 @@ def single_model_finetune( self.wrapper.share_params( shared_links, resume=(resuming and not self.finetune_update_stat) or self.rank != 0, - model_key_prob_map=dict(zip(self.model_keys, self.model_prob, strict=False)), + model_key_prob_map=dict(zip(self.model_keys, self.model_prob)), data_stat_protect=_data_stat_protect[0], ) diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py index ebab848564..c991e59daa 100644 --- a/deepmd/pt/utils/dataloader.py +++ b/deepmd/pt/utils/dataloader.py @@ -157,7 +157,7 @@ def construct_dataset(system: str) -> DeepmdDataSetForLoader: else: self.batch_sizes = batch_size * np.ones(len(systems), dtype=int) assert len(self.systems) == len(self.batch_sizes) - for system, batch_size in zip(self.systems, self.batch_sizes, strict=False): + for system, batch_size in zip(self.systems, self.batch_sizes): if dist.is_available() and dist.is_initialized(): system_sampler = DistributedSampler(system) self.sampler_list.append(system_sampler) diff --git a/deepmd/pt/utils/nlist.py b/deepmd/pt/utils/nlist.py index 21b678ca1b..7f74e65f26 100644 --- a/deepmd/pt/utils/nlist.py +++ b/deepmd/pt/utils/nlist.py @@ -398,7 +398,7 @@ def build_multiple_neighbor_list( rr.masked_fill(nlist_mask, float("inf")) nlist0 = nlist ret = {} - for rc, ns in zip(rcuts[::-1], nsels[::-1], strict=False): + for rc, ns in zip(rcuts[::-1], nsels[::-1]): nlist0 = nlist0[:, :, :ns].masked_fill(rr[:, :, :ns] > rc, -1) ret[get_multiple_nlist_key(rc, ns)] = nlist0 return ret From d54e691f30ed5db622c20ddbde5bc66420119b44 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 8 Jan 2026 19:10:58 +0800 Subject: [PATCH 06/12] exclude pd --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index e80089e6b8..bc7027de0f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -432,7 +432,7 @@ runtime-evaluated-base-classes = ["torch.nn.Module"] "deepmd/tf/**" = ["TID253", "ANN"] "deepmd/pt/**" = ["TID253", "B905"] "deepmd/jax/**" = ["TID253"] -"deepmd/pd/**" = ["TID253", "ANN"] +"deepmd/pd/**" = ["TID253", "ANN", "B905"] "source/**" = ["ANN"] "source/tests/tf/**" = ["TID253", "ANN"] From 459b7458a64fb7a8fc55e5d44e634991b887c01f Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 8 Jan 2026 19:11:21 +0800 Subject: [PATCH 07/12] revert pd --- deepmd/pd/infer/deep_eval.py | 1 - deepmd/pd/model/descriptor/se_a.py | 7 +------ deepmd/pd/model/model/transform_output.py | 2 +- deepmd/pd/utils/dataloader.py | 2 +- deepmd/pd/utils/dataset.py | 5 ++++- deepmd/pd/utils/nlist.py | 2 +- 6 files changed, 8 insertions(+), 11 deletions(-) diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index fcf135741b..67c435ab3f 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -385,7 +385,6 @@ def eval( zip( [x.name for x in request_defs], out, - strict=False, ) ) diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py index 7d626425ec..17ef5d67c4 100644 --- a/deepmd/pd/model/descriptor/se_a.py +++ b/deepmd/pd/model/descriptor/se_a.py @@ -763,12 +763,7 @@ def forward( # nfnl x nnei exclude_mask = self.emask(nlist, extended_atype).reshape([nfnl, self.nnei]) for embedding_idx, (ll, compress_data_ii, compress_info_ii) in enumerate( - zip( - self.filter_layers.networks, - self.compress_data, - self.compress_info, - strict=False, - ) + zip(self.filter_layers.networks, self.compress_data, self.compress_info) ): if self.type_one_side: ii = embedding_idx diff --git a/deepmd/pd/model/model/transform_output.py b/deepmd/pd/model/model/transform_output.py index 692e01a40c..47004265c7 100644 --- a/deepmd/pd/model/model/transform_output.py +++ b/deepmd/pd/model/model/transform_output.py @@ -122,7 +122,7 @@ def take_deriv( split_vv1 = paddle.split(vv1, [1] * size, axis=-1) split_svv1 = paddle.split(svv1, [1] * size, axis=-1) split_ff, split_avir = [], [] - for vvi, svvi in zip(split_vv1, split_svv1, strict=False): + for vvi, svvi in zip(split_vv1, split_svv1): # nf x nloc x 3, nf x nloc x 9 ffi, aviri = task_deriv_one( vvi, diff --git a/deepmd/pd/utils/dataloader.py b/deepmd/pd/utils/dataloader.py index cf36ea2d05..0cb8adbc63 100644 --- a/deepmd/pd/utils/dataloader.py +++ b/deepmd/pd/utils/dataloader.py @@ -167,7 +167,7 @@ def construct_dataset(system): else: self.batch_sizes = batch_size * np.ones(len(systems), dtype=int) assert len(self.systems) == len(self.batch_sizes) - for system, batch_size in zip(self.systems, self.batch_sizes, strict=False): + for system, batch_size in zip(self.systems, self.batch_sizes): if dist.is_available() and dist.is_initialized(): system_batch_sampler = DistributedBatchSampler( system, diff --git a/deepmd/pd/utils/dataset.py b/deepmd/pd/utils/dataset.py index fa9106044c..5accd0315b 100644 --- a/deepmd/pd/utils/dataset.py +++ b/deepmd/pd/utils/dataset.py @@ -5,6 +5,9 @@ Dataset, ) +from deepmd.pd.utils.env import ( + NUM_WORKERS, +) from deepmd.utils.data import ( DataRequirementItem, DeepmdData, @@ -32,7 +35,7 @@ def __len__(self): def __getitem__(self, index): """Get a frame from the selected system.""" - b_data = self._data_system.get_item_paddle(index) + b_data = self._data_system.get_item_paddle(index, max(1, NUM_WORKERS)) b_data["natoms"] = self._natoms_vec return b_data diff --git a/deepmd/pd/utils/nlist.py b/deepmd/pd/utils/nlist.py index 203bdff46f..cd041a345c 100644 --- a/deepmd/pd/utils/nlist.py +++ b/deepmd/pd/utils/nlist.py @@ -404,7 +404,7 @@ def build_multiple_neighbor_list( rr.masked_fill(nlist_mask, float("inf")) nlist0 = nlist ret = {} - for rc, ns in zip(rcuts[::-1], nsels[::-1], strict=False): + for rc, ns in zip(rcuts[::-1], nsels[::-1]): nlist0 = nlist0[:, :, :ns].masked_fill(rr[:, :, :ns] > rc, -1) ret[get_multiple_nlist_key(rc, ns)] = nlist0 return ret From c39d25b1eb2c088c78c251b1c074d3d138602448 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 8 Jan 2026 19:13:59 +0800 Subject: [PATCH 08/12] try strict=True --- deepmd/calculator.py | 2 +- .../atomic_model/linear_atomic_model.py | 6 ++--- deepmd/dpmodel/descriptor/hybrid.py | 4 ++-- deepmd/dpmodel/infer/deep_eval.py | 2 +- deepmd/dpmodel/utils/nlist.py | 2 +- deepmd/jax/infer/deep_eval.py | 2 +- deepmd/tf/descriptor/loc_frame.py | 2 +- deepmd/tf/descriptor/se_a.py | 4 ++-- deepmd/tf/descriptor/se_a_ef.py | 2 +- deepmd/tf/descriptor/se_atten.py | 4 ++-- deepmd/tf/descriptor/se_r.py | 2 +- deepmd/tf/descriptor/se_t.py | 2 +- deepmd/tf/infer/deep_eval.py | 2 +- deepmd/tf/nvnmd/entrypoints/mapt.py | 10 ++++---- deepmd/tf/train/trainer.py | 2 +- deepmd/utils/batch_size.py | 2 +- deepmd/utils/data_system.py | 2 +- deepmd/utils/model_branch_dict.py | 4 ++-- deepmd/utils/update_sel.py | 2 +- .../dpmodel/test_pairtab_atomic_model.py | 2 +- source/tests/consistent/common.py | 24 +++++++++---------- .../tests/consistent/fitting/test_dipole.py | 2 +- source/tests/consistent/io/test_io.py | 4 ++-- source/tests/tf/test_linear_model.py | 2 +- source/tests/tf/test_model_devi.py | 4 ++-- source/tests/tf/test_model_devi_mix.py | 2 +- source/tests/tf/test_nvnmd_entrypoints.py | 4 ++-- 27 files changed, 51 insertions(+), 51 deletions(-) diff --git a/deepmd/calculator.py b/deepmd/calculator.py index 1b40cc09d6..1d8e955de7 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -103,7 +103,7 @@ def __init__( self.type_dict = type_dict else: self.type_dict = dict( - zip(self.dp.get_type_map(), range(self.dp.get_ntypes()), strict=False) + zip(self.dp.get_type_map(), range(self.dp.get_ntypes()), strict=True) ) def calculate( diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 766f8a1821..380c300216 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -155,7 +155,7 @@ def get_model_sels(self) -> list[int | list[int]]: def _sort_rcuts_sels(self) -> tuple[tuple[Array, Array], list[int]]: # sort the pair of rcut and sels in ascending order, first based on sel, then on rcut. zipped = sorted( - zip(self.get_model_rcuts(), self.get_model_nsels(), strict=False), + zip(self.get_model_rcuts(), self.get_model_nsels(), strict=True), key=lambda x: (x[1], x[0]), ) return [p[0] for p in zipped], [p[1] for p in zipped] @@ -236,13 +236,13 @@ def forward_atomic( raw_nlists = [ nlists[get_multiple_nlist_key(rcut, sel)] for rcut, sel in zip( - self.get_model_rcuts(), self.get_model_nsels(), strict=False + self.get_model_rcuts(), self.get_model_nsels(), strict=True ) ] nlists_ = [ nl if mt else nlist_distinguish_types(nl, extended_atype, sel) for mt, nl, sel in zip( - self.mixed_types_list, raw_nlists, self.get_model_sels(), strict=False + self.mixed_types_list, raw_nlists, self.get_model_sels(), strict=True ) ] ener_list = [] diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 95d4716d5c..2cb8585d77 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -101,7 +101,7 @@ def __init__( start_idx = np.cumsum(np.pad(hybrid_sel, (1, 0), "constant"))[:-1] end_idx = start_idx + np.array(sub_sel) cut_idx = np.concatenate( - [range(ss, ee) for ss, ee in zip(start_idx, end_idx, strict=False)] + [range(ss, ee) for ss, ee in zip(start_idx, end_idx, strict=True)] ) nlist_cut_idx.append(cut_idx) self.nlist_cut_idx = nlist_cut_idx @@ -310,7 +310,7 @@ def call( ) else: nl_distinguish_types = None - for descrpt, nci in zip(self.descrpt_list, self.nlist_cut_idx, strict=False): + for descrpt, nci in zip(self.descrpt_list, self.nlist_cut_idx, strict=True): # cut the nlist to the correct length if self.mixed_types() == descrpt.mixed_types(): nl = xp.take(nlist, nci, axis=2) diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index bb366a84ca..8088ba1d2f 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -229,7 +229,7 @@ def eval( zip( [x.name for x in request_defs], out, - strict=False, + strict=True, ) ) diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index 1e361589c4..7dd41f45bd 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -234,7 +234,7 @@ def build_multiple_neighbor_list( rr = xp.where(nlist_mask, xp.full_like(rr, float("inf")), rr) nlist0 = nlist ret = {} - for rc, ns in zip(rcuts[::-1], nsels[::-1], strict=False): + for rc, ns in zip(rcuts[::-1], nsels[::-1], strict=True): tnlist_1 = nlist0[:, :, :ns] tnlist_1 = xp.where(rr[:, :, :ns] > rc, xp.full_like(tnlist_1, -1), tnlist_1) ret[get_multiple_nlist_key(rc, ns)] = tnlist_1 diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 8a2d5efef6..7a7723b181 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -249,7 +249,7 @@ def eval( zip( [x.name for x in request_defs], out, - strict=False, + strict=True, ) ) diff --git a/deepmd/tf/descriptor/loc_frame.py b/deepmd/tf/descriptor/loc_frame.py index 1fd7f521e5..187f2a5e31 100644 --- a/deepmd/tf/descriptor/loc_frame.py +++ b/deepmd/tf/descriptor/loc_frame.py @@ -187,7 +187,7 @@ def compute_input_stats( sumn = [] sumv2 = [] for cc, bb, tt, nn, mm in zip( - data_coord, data_box, data_atype, natoms_vec, mesh, strict=False + data_coord, data_box, data_atype, natoms_vec, mesh, strict=True ): sysv, sysv2, sysn = self._compute_dstats_sys_nonsmth(cc, bb, tt, nn, mm) sumv.append(sysv) diff --git a/deepmd/tf/descriptor/se_a.py b/deepmd/tf/descriptor/se_a.py index 52b38bbcad..5feea7fd63 100644 --- a/deepmd/tf/descriptor/se_a.py +++ b/deepmd/tf/descriptor/se_a.py @@ -374,7 +374,7 @@ def compute_input_stats( sumr2 = [] suma2 = [] for cc, bb, tt, nn, mm in zip( - data_coord, data_box, data_atype, natoms_vec, mesh, strict=False + data_coord, data_box, data_atype, natoms_vec, mesh, strict=True ): sysr, sysr2, sysa, sysa2, sysn = self._compute_dstats_sys_smth( cc, bb, tt, nn, mm @@ -1331,7 +1331,7 @@ def init_variables( start_index_old[0] = 0 for nn, oo, ii, jj in zip( - n_descpt, n_descpt_old, start_index, start_index_old, strict=False + n_descpt, n_descpt_old, start_index, start_index_old, strict=True ): if nn < oo: # new size is smaller, copy part of std diff --git a/deepmd/tf/descriptor/se_a_ef.py b/deepmd/tf/descriptor/se_a_ef.py index 146ffc424d..b2e161938d 100644 --- a/deepmd/tf/descriptor/se_a_ef.py +++ b/deepmd/tf/descriptor/se_a_ef.py @@ -425,7 +425,7 @@ def compute_input_stats( natoms_vec, mesh, data_efield, - strict=False, + strict=True, ): sysr, sysr2, sysa, sysa2, sysn = self._compute_dstats_sys_smth( cc, bb, tt, nn, mm, ee diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index 8cce8f7da7..dc30c4088e 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -385,7 +385,7 @@ def compute_input_stats( natoms_vec, mesh, real_natoms_vec, - strict=False, + strict=True, ): sysr, sysr2, sysa, sysa2, sysn = self._compute_dstats_sys_smth( cc, bb, tt, nn, mm, mixed_type, r_n @@ -398,7 +398,7 @@ def compute_input_stats( suma2.append(sysa2) else: for cc, bb, tt, nn, mm in zip( - data_coord, data_box, data_atype, natoms_vec, mesh, strict=False + data_coord, data_box, data_atype, natoms_vec, mesh, strict=True ): sysr, sysr2, sysa, sysa2, sysn = self._compute_dstats_sys_smth( cc, bb, tt, nn, mm diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py index 9bf553326f..07452f9ccd 100644 --- a/deepmd/tf/descriptor/se_r.py +++ b/deepmd/tf/descriptor/se_r.py @@ -271,7 +271,7 @@ def compute_input_stats( sumn = [] sumr2 = [] for cc, bb, tt, nn, mm in zip( - data_coord, data_box, data_atype, natoms_vec, mesh, strict=False + data_coord, data_box, data_atype, natoms_vec, mesh, strict=True ): sysr, sysr2, sysn = self._compute_dstats_sys_se_r(cc, bb, tt, nn, mm) sumr.append(sysr) diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index 96fc5d9189..bb7b64bf5f 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -258,7 +258,7 @@ def compute_input_stats( sumr2 = [] suma2 = [] for cc, bb, tt, nn, mm in zip( - data_coord, data_box, data_atype, natoms_vec, mesh, strict=False + data_coord, data_box, data_atype, natoms_vec, mesh, strict=True ): sysr, sysr2, sysa, sysa2, sysn = self._compute_dstats_sys_smth( cc, bb, tt, nn, mm diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index a7eedb2e16..b8c029a1d4 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -761,7 +761,7 @@ def eval( output_dict = { odef.name: oo - for oo, odef in zip(output, self.output_def.var_defs.values(), strict=False) + for oo, odef in zip(output, self.output_def.var_defs.values(), strict=True) } # ugly!! if self.modifier_type is not None and issubclass(self.model_type, DeepPot): diff --git a/deepmd/tf/nvnmd/entrypoints/mapt.py b/deepmd/tf/nvnmd/entrypoints/mapt.py index 7e48bde300..977069058a 100644 --- a/deepmd/tf/nvnmd/entrypoints/mapt.py +++ b/deepmd/tf/nvnmd/entrypoints/mapt.py @@ -439,11 +439,11 @@ def run_u2s(self): u = N2 * np.reshape(np.arange(0, N + 1) / N, [-1, 1]) # pylint: disable=no-explicit-dtype res_lst = run_sess(sess, vals, feed_dict={dic_ph["u"]: u}) - res_dic = dict(zip(keys, res_lst, strict=False)) + res_dic = dict(zip(keys, res_lst, strict=True)) u2 = N2 * np.reshape(np.arange(0, N * 16 + 1) / (N * 16), [-1, 1]) # pylint: disable=no-explicit-dtype res_lst2 = run_sess(sess, vals, feed_dict={dic_ph["u"]: u2}) - res_dic2 = dict(zip(keys, res_lst2, strict=False)) # reference for compare + res_dic2 = dict(zip(keys, res_lst2, strict=True)) # reference for compare # change value for tt in range(ndim): @@ -536,11 +536,11 @@ def run_s2g(self): s = N2 * np.reshape(np.arange(0, N + 1) / N, [-1, 1]) + smin_ # pylint: disable=no-explicit-dtype res_lst = run_sess(sess, vals, feed_dict={dic_ph["s"]: s}) - res_dic = dict(zip(keys, res_lst, strict=False)) + res_dic = dict(zip(keys, res_lst, strict=True)) s2 = N2 * np.reshape(np.arange(0, N * 16 + 1) / (N * 16), [-1, 1]) + smin_ # pylint: disable=no-explicit-dtype res_lst2 = run_sess(sess, vals, feed_dict={dic_ph["s"]: s2}) - res_dic2 = dict(zip(keys, res_lst2, strict=False)) + res_dic2 = dict(zip(keys, res_lst2, strict=True)) sess.close() return res_dic, res_dic2 @@ -601,7 +601,7 @@ def run_t2g(self): vals = list(dic_ph.values()) # res_lst = run_sess(sess, vals, feed_dict={}) - res_dic = dict(zip(keys, res_lst, strict=False)) + res_dic = dict(zip(keys, res_lst, strict=True)) sess.close() return res_dic diff --git a/deepmd/tf/train/trainer.py b/deepmd/tf/train/trainer.py index e10cff967a..4af59fd290 100644 --- a/deepmd/tf/train/trainer.py +++ b/deepmd/tf/train/trainer.py @@ -967,4 +967,4 @@ def get_data_dict(self, batch_list: list[np.ndarray]) -> dict[str, np.ndarray]: dict[str, np.ndarray] The dict of the loaded data. """ - return dict(zip(self.data_keys, batch_list, strict=False)) + return dict(zip(self.data_keys, batch_list, strict=True)) diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py index effd3bd53f..e701e82ec6 100644 --- a/deepmd/utils/batch_size.py +++ b/deepmd/utils/batch_size.py @@ -248,7 +248,7 @@ def concate_result(r: list[Any]) -> Any: return ret if not returned_dict: - r_list = [concate_result(r) for r in zip(*results, strict=False)] + r_list = [concate_result(r) for r in zip(*results, strict=True)] r = tuple(r_list) if len(r) == 1: # avoid returning tuple if callable doesn't return tuple diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 6c4fe4a190..c9edfc2298 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -173,7 +173,7 @@ def __init__( filtered_data_systems = [] filtered_system_dirs = [] for sys_dir, data_sys in zip( - self.system_dirs, self.data_systems, strict=False + self.system_dirs, self.data_systems, strict=True ): if data_sys.get_natoms() <= rule: filtered_data_systems.append(data_sys) diff --git a/deepmd/utils/model_branch_dict.py b/deepmd/utils/model_branch_dict.py index 01348c3cee..ccf085cb2b 100644 --- a/deepmd/utils/model_branch_dict.py +++ b/deepmd/utils/model_branch_dict.py @@ -169,7 +169,7 @@ def as_table(self) -> str: # Step 3: Determine actual width for each column # For the first two columns, we already decided the exact widths above. col_widths: list[int] = [] - for idx, col in enumerate(zip(*wrapped_rows, strict=False)): + for idx, col in enumerate(zip(*wrapped_rows, strict=True)): if idx == 0: col_widths.append(branch_col_width) elif idx == 1: @@ -188,7 +188,7 @@ def draw_row_line(cells_parts: list[list[str]]) -> str: "| " + " | ".join( part.ljust(width) - for part, width in zip(cells_parts, col_widths, strict=False) + for part, width in zip(cells_parts, col_widths, strict=True) ) + " |" ) diff --git a/deepmd/utils/update_sel.py b/deepmd/utils/update_sel.py index 75e62f5006..7b93a8501d 100644 --- a/deepmd/utils/update_sel.py +++ b/deepmd/utils/update_sel.py @@ -43,7 +43,7 @@ def update_one_sel( sel = [int(self.wrap_up_4(ii * ratio)) for ii in tmp_sel] else: # sel is set by user - for ii, (tt, dd) in enumerate(zip(tmp_sel, sel, strict=False)): + for ii, (tt, dd) in enumerate(zip(tmp_sel, sel, strict=True)): if dd and tt > dd: # we may skip warning for sel=0, where the user is likely # to exclude such type in the descriptor diff --git a/source/tests/common/dpmodel/test_pairtab_atomic_model.py b/source/tests/common/dpmodel/test_pairtab_atomic_model.py index 19cf21d13c..1e87bb4bd7 100644 --- a/source/tests/common/dpmodel/test_pairtab_atomic_model.py +++ b/source/tests/common/dpmodel/test_pairtab_atomic_model.py @@ -158,7 +158,7 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: 0.035, 0.025, ], - strict=False, + strict=True, ): extended_coord = np.array( [ diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 42814566f3..88fad4e10b 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -355,7 +355,7 @@ def test_tf_consistent_with_ref(self) -> None: data2.pop("@version") np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): np.testing.assert_allclose( rr1.ravel(), rr2.ravel(), rtol=self.rtol, atol=self.atol ) @@ -372,7 +372,7 @@ def test_tf_self_consistent(self) -> None: obj2 = self.tf_class.deserialize(data1, suffix=self.unique_id) ret2, data2 = self.get_tf_ret_serialization_from_cls(obj2) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -391,7 +391,7 @@ def test_dp_consistent_with_ref(self) -> None: ret2 = self.extract_ret(ret2, self.RefBackend.DP) data2 = dp_obj.serialize() np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): if rr1 is SKIP_FLAG or rr2 is SKIP_FLAG: continue np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) @@ -407,7 +407,7 @@ def test_dp_self_consistent(self) -> None: obj1 = self.dp_class.deserialize(data1) ret2, data2 = self.get_dp_ret_serialization_from_cls(obj1) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): if isinstance(rr1, np.ndarray) and isinstance(rr2, np.ndarray): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -436,7 +436,7 @@ def test_pt_consistent_with_ref(self) -> None: data1.pop("@variables", None) data2.pop("@variables", None) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -449,7 +449,7 @@ def test_pt_self_consistent(self) -> None: obj2 = self.pt_class.deserialize(data1) ret2, data2 = self.get_pt_ret_serialization_from_cls(obj2) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): if isinstance(rr1, np.ndarray) and isinstance(rr2, np.ndarray): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -473,7 +473,7 @@ def test_jax_consistent_with_ref(self) -> None: data1.pop("@variables", None) data2.pop("@variables", None) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -486,7 +486,7 @@ def test_jax_self_consistent(self) -> None: obj1 = self.jax_class.deserialize(data1) ret2, data2 = self.get_jax_ret_serialization_from_cls(obj1) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): if isinstance(rr1, np.ndarray) and isinstance(rr2, np.ndarray): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -515,7 +515,7 @@ def test_pd_consistent_with_ref(self): data1.pop("@variables", None) data2.pop("@variables", None) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -528,7 +528,7 @@ def test_pd_self_consistent(self): obj2 = self.pd_class.deserialize(data1) ret2, data2 = self.get_pd_ret_serialization_from_cls(obj2) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): if isinstance(rr1, np.ndarray) and isinstance(rr2, np.ndarray): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -550,7 +550,7 @@ def test_array_api_strict_consistent_with_ref(self) -> None: ret2 = self.extract_ret(ret2, self.RefBackend.ARRAY_API_STRICT) data2 = array_api_strict_obj.serialize() np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @@ -564,7 +564,7 @@ def test_array_api_strict_self_consistent(self) -> None: obj1 = self.array_api_strict_class.deserialize(data1) ret2, data2 = self.get_array_api_strict_ret_serialization_from_cls(obj1) np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): if isinstance(rr1, np.ndarray) and isinstance(rr2, np.ndarray): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" diff --git a/source/tests/consistent/fitting/test_dipole.py b/source/tests/consistent/fitting/test_dipole.py index bf2907df03..a1c48d102d 100644 --- a/source/tests/consistent/fitting/test_dipole.py +++ b/source/tests/consistent/fitting/test_dipole.py @@ -285,7 +285,7 @@ def test_tf_consistent_with_ref(self) -> None: if networks is None: data1["nets"]["networks"][ii] = None np.testing.assert_equal(data1, data2) - for rr1, rr2 in zip(ret1, ret2, strict=False): + for rr1, rr2 in zip(ret1, ret2, strict=True): np.testing.assert_allclose( rr1.ravel()[: rr2.size], rr2.ravel(), rtol=self.rtol, atol=self.atol ) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index cc7b3dcde1..338fc22561 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -202,14 +202,14 @@ def test_deep_eval(self) -> None: ) rets_nopbc.append(ret) for ret in rets[1:]: - for vv1, vv2 in zip(rets[0], ret, strict=False): + for vv1, vv2 in zip(rets[0], ret, strict=True): if np.isnan(vv2).all(): # expect all nan if not supported continue np.testing.assert_allclose(vv1, vv2, rtol=1e-12, atol=1e-12) for idx, ret in enumerate(rets_nopbc[1:]): - for vv1, vv2 in zip(rets_nopbc[0], ret, strict=False): + for vv1, vv2 in zip(rets_nopbc[0], ret, strict=True): if np.isnan(vv2).all(): # expect all nan if not supported continue diff --git a/source/tests/tf/test_linear_model.py b/source/tests/tf/test_linear_model.py index 5746eab84c..6743c4658b 100644 --- a/source/tests/tf/test_linear_model.py +++ b/source/tests/tf/test_linear_model.py @@ -37,7 +37,7 @@ def setUp(self) -> None: os.path.join(infer_path, "deeppot-1.pbtxt"), ] self.graph_dirs = [pbtxt.replace("pbtxt", "pb") for pbtxt in self.pbtxts] - for pbtxt, pb in zip(self.pbtxts, self.graph_dirs, strict=False): + for pbtxt, pb in zip(self.pbtxts, self.graph_dirs, strict=True): convert_pbtxt_to_pb(pbtxt, pb) self.graphs = [DeepPotential(pb) for pb in self.graph_dirs] diff --git a/source/tests/tf/test_model_devi.py b/source/tests/tf/test_model_devi.py index de02dedba5..e213c43893 100644 --- a/source/tests/tf/test_model_devi.py +++ b/source/tests/tf/test_model_devi.py @@ -41,7 +41,7 @@ def setUp(self) -> None: os.path.join(infer_path, "deeppot-1.pbtxt"), ] self.graph_dirs = [pbtxt.replace("pbtxt", "pb") for pbtxt in self.pbtxts] - for pbtxt, pb in zip(self.pbtxts, self.graph_dirs, strict=False): + for pbtxt, pb in zip(self.pbtxts, self.graph_dirs, strict=True): convert_pbtxt_to_pb(pbtxt, pb) self.graphs = [DeepPotential(pb) for pb in self.graph_dirs] self.output = os.path.join(tests_path, "model_devi.out") @@ -216,7 +216,7 @@ def setUpClass(cls) -> None: os.path.join(infer_path, "fparam_aparam.pbtxt"), ] cls.graph_dirs = [pbtxt.replace("pbtxt", "pb") for pbtxt in cls.pbtxts] - for pbtxt, pb in zip(cls.pbtxts, cls.graph_dirs, strict=False): + for pbtxt, pb in zip(cls.pbtxts, cls.graph_dirs, strict=True): convert_pbtxt_to_pb(pbtxt, pb) cls.graphs = [DeepPotential(pb) for pb in cls.graph_dirs] diff --git a/source/tests/tf/test_model_devi_mix.py b/source/tests/tf/test_model_devi_mix.py index b3907751aa..2c660d7a22 100644 --- a/source/tests/tf/test_model_devi_mix.py +++ b/source/tests/tf/test_model_devi_mix.py @@ -58,7 +58,7 @@ def setUp(self) -> None: os.path.join(infer_path, "se_atten_no_atten_2.pbtxt"), ] self.graph_dirs = [pbtxt.replace("pbtxt", "pb") for pbtxt in self.pbtxts] - for pbtxt, pb in zip(self.pbtxts, self.graph_dirs, strict=False): + for pbtxt, pb in zip(self.pbtxts, self.graph_dirs, strict=True): convert_pbtxt_to_pb(pbtxt, pb) self.graphs = [DeepPotential(pb) for pb in self.graph_dirs] self.output = os.path.join(tests_path, "model_devi.out") diff --git a/source/tests/tf/test_nvnmd_entrypoints.py b/source/tests/tf/test_nvnmd_entrypoints.py index 2e3c3f2f37..4a6877761e 100644 --- a/source/tests/tf/test_nvnmd_entrypoints.py +++ b/source/tests/tf/test_nvnmd_entrypoints.py @@ -466,7 +466,7 @@ def test_model_qnn_v0(self) -> None: tensordic[key] = graph.get_tensor_by_name(key + ":0") # get value valuelist = sess.run(list(tensordic.values()), feed_dict=feed_dict) - valuedic = dict(zip(tensordic.keys(), valuelist, strict=False)) + valuedic = dict(zip(tensordic.keys(), valuelist, strict=True)) # test # o_descriptor idx = np.array([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]) @@ -783,7 +783,7 @@ def test_model_qnn_v1(self) -> None: tensordic[key] = graph.get_tensor_by_name(key + ":0") # get value valuelist = sess.run(list(tensordic.values()), feed_dict=feed_dict) - valuedic = dict(zip(tensordic.keys(), valuelist, strict=False)) + valuedic = dict(zip(tensordic.keys(), valuelist, strict=True)) # test # o_descriptor idx = np.array([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]) From de735a7a8398f571c3ad95f303670f6b808cda8c Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 10 Jan 2026 15:24:10 +0800 Subject: [PATCH 09/12] Refactor IOTest to separate atomic results in evaluation tests --- source/tests/consistent/io/test_io.py | 37 +++++++++++++++------------ 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 338fc22561..982d56d8fa 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -140,7 +140,9 @@ def test_deep_eval(self) -> None: nframes = self.atype.shape[0] prefix = "test_consistent_io_" + self.__class__.__name__.lower() rets = [] + rets_atomic = [] rets_nopbc = [] + rets_nopbc_atomic = [] for backend_name, suffix_idx in ( # unfortunately, jax2tf cannot work with tf v1 behaviors ("jax", 2) if DP_TEST_TF2_ONLY else ("tensorflow", 0), @@ -183,7 +185,8 @@ def test_deep_eval(self) -> None: aparam=aparam, atomic=True, ) - rets.append(ret) + rets.append(ret[:3]) + rets_atomic.append(ret[3:]) ret = deep_eval.eval( self.coords, None, @@ -200,22 +203,24 @@ def test_deep_eval(self) -> None: aparam=aparam, atomic=True, ) - rets_nopbc.append(ret) - for ret in rets[1:]: - for vv1, vv2 in zip(rets[0], ret, strict=True): - if np.isnan(vv2).all(): - # expect all nan if not supported - continue - np.testing.assert_allclose(vv1, vv2, rtol=1e-12, atol=1e-12) + rets_nopbc.append(ret[:3]) + rets_nopbc_atomic.append(ret[3:]) - for idx, ret in enumerate(rets_nopbc[1:]): - for vv1, vv2 in zip(rets_nopbc[0], ret, strict=True): - if np.isnan(vv2).all(): - # expect all nan if not supported - continue - np.testing.assert_allclose( - vv1, vv2, rtol=1e-12, atol=1e-12, err_msg=f"backend {idx + 1}" - ) + for rets_idx, rets_x in enumerate( + (rets, rets_atomic, rets_nopbc, rets_nopbc_atomic) + ): + for idx, ret in enumerate(rets_x[1:]): + for vv1, vv2 in zip(rets_x[0], ret, strict=True): + if np.isnan(vv2).all(): + # expect all nan if not supported + continue + np.testing.assert_allclose( + vv1, + vv2, + rtol=1e-12, + atol=1e-12, + err_msg=f"backend {idx + 1} for rets_idx {rets_idx}", + ) class TestDeepPot(unittest.TestCase, IOTest): From 69189dbd72539ba278986c2eb3f7dac2dac02fb6 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 10 Jan 2026 23:51:49 +0800 Subject: [PATCH 10/12] Fix strict mode handling in sel update for TF spin models --- deepmd/utils/update_sel.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/deepmd/utils/update_sel.py b/deepmd/utils/update_sel.py index 7b93a8501d..589b17edac 100644 --- a/deepmd/utils/update_sel.py +++ b/deepmd/utils/update_sel.py @@ -43,7 +43,9 @@ def update_one_sel( sel = [int(self.wrap_up_4(ii * ratio)) for ii in tmp_sel] else: # sel is set by user - for ii, (tt, dd) in enumerate(zip(tmp_sel, sel, strict=True)): + # TODO: Fix len(tmp_sel) != len(sel) for TF spin models when strict is True + # error reported by source/tests/tf/test_init_frz_model_spin.py + for ii, (tt, dd) in enumerate(zip(tmp_sel, sel, strict=False)): if dd and tt > dd: # we may skip warning for sel=0, where the user is likely # to exclude such type in the descriptor From 3e628a943c72124e8948e312a470323dcba310a1 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 11 Jan 2026 00:48:09 +0800 Subject: [PATCH 11/12] fix(tests): enable strict mode in zip assertions across multiple test files --- source/tests/pd/model/test_nlist.py | 2 +- source/tests/pd/model/test_se_e2_a.py | 2 +- source/tests/pd/test_utils.py | 2 +- source/tests/pt/model/test_descriptor_se_r.py | 2 +- source/tests/pt/model/test_nlist.py | 2 +- source/tests/pt/model/test_pairtab_atomic_model.py | 2 +- source/tests/pt/model/test_se_e2_a.py | 2 +- source/tests/pt/model/test_se_t.py | 2 +- source/tests/pt/test_utils.py | 2 +- 9 files changed, 9 insertions(+), 9 deletions(-) diff --git a/source/tests/pd/model/test_nlist.py b/source/tests/pd/model/test_nlist.py index b45767aa52..09713c412b 100644 --- a/source/tests/pd/model/test_nlist.py +++ b/source/tests/pd/model/test_nlist.py @@ -254,7 +254,7 @@ def test_build_directional_nlist(self): self.coord, self.atype, self.cell, self.rcut ) for distinguish_types, mysel in zip( - [True, False], [sum(self.nsel), 300], strict=False + [True, False], [sum(self.nsel), 300], strict=True ): # full neighbor list nlist_full = build_neighbor_list( diff --git a/source/tests/pd/model/test_se_e2_a.py b/source/tests/pd/model/test_se_e2_a.py index c7809586e6..f512f9e1a1 100644 --- a/source/tests/pd/model/test_se_e2_a.py +++ b/source/tests/pd/model/test_se_e2_a.py @@ -96,7 +96,7 @@ def test_consistency( self.atype_ext, self.nlist, ) - for aa, bb in zip([rd1, gr1, sw1], [rd2, gr2, sw2], strict=False): + for aa, bb in zip([rd1, gr1, sw1], [rd2, gr2, sw2], strict=True): np.testing.assert_allclose( aa.detach().cpu().numpy(), bb, diff --git a/source/tests/pd/test_utils.py b/source/tests/pd/test_utils.py index f4aafe97af..39e9dcfe43 100644 --- a/source/tests/pd/test_utils.py +++ b/source/tests/pd/test_utils.py @@ -21,7 +21,7 @@ def test_to_numpy(self): for ptp, npp in zip( [paddle.float16, paddle.float32, paddle.float64], [np.float16, np.float32, np.float64], - strict=False, + strict=True, ): foo = foo.astype(npp) bar = to_paddle_tensor(foo) diff --git a/source/tests/pt/model/test_descriptor_se_r.py b/source/tests/pt/model/test_descriptor_se_r.py index 86722bc19b..64811f774c 100644 --- a/source/tests/pt/model/test_descriptor_se_r.py +++ b/source/tests/pt/model/test_descriptor_se_r.py @@ -100,7 +100,7 @@ def test_consistency( self.atype_ext, self.nlist, ) - for aa, bb in zip([rd1, sw1], [rd2, sw2], strict=False): + for aa, bb in zip([rd1, sw1], [rd2, sw2], strict=True): np.testing.assert_allclose( aa.detach().cpu().numpy(), bb, diff --git a/source/tests/pt/model/test_nlist.py b/source/tests/pt/model/test_nlist.py index cfe63d0fa7..cac5e5294a 100644 --- a/source/tests/pt/model/test_nlist.py +++ b/source/tests/pt/model/test_nlist.py @@ -246,7 +246,7 @@ def test_build_directional_nlist(self) -> None: self.coord, self.atype, self.cell, self.rcut ) for distinguish_types, mysel in zip( - [True, False], [sum(self.nsel), 300], strict=False + [True, False], [sum(self.nsel), 300], strict=True ): # full neighbor list nlist_full = build_neighbor_list( diff --git a/source/tests/pt/model/test_pairtab_atomic_model.py b/source/tests/pt/model/test_pairtab_atomic_model.py index bd9a629a10..0f324cbf51 100644 --- a/source/tests/pt/model/test_pairtab_atomic_model.py +++ b/source/tests/pt/model/test_pairtab_atomic_model.py @@ -219,7 +219,7 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: 0.035, 0.025, ], - strict=False, + strict=True, ): extended_coord = torch.tensor( [ diff --git a/source/tests/pt/model/test_se_e2_a.py b/source/tests/pt/model/test_se_e2_a.py index d87f5c8348..bdf7e3b4b9 100644 --- a/source/tests/pt/model/test_se_e2_a.py +++ b/source/tests/pt/model/test_se_e2_a.py @@ -96,7 +96,7 @@ def test_consistency( self.atype_ext, self.nlist, ) - for aa, bb in zip([rd1, gr1, sw1], [rd2, gr2, sw2], strict=False): + for aa, bb in zip([rd1, gr1, sw1], [rd2, gr2, sw2], strict=True): np.testing.assert_allclose( aa.detach().cpu().numpy(), bb, diff --git a/source/tests/pt/model/test_se_t.py b/source/tests/pt/model/test_se_t.py index cddb455295..8f67d4a0aa 100644 --- a/source/tests/pt/model/test_se_t.py +++ b/source/tests/pt/model/test_se_t.py @@ -100,7 +100,7 @@ def test_consistency( self.atype_ext, self.nlist, ) - for aa, bb in zip([rd1, sw1], [rd2, sw2], strict=False): + for aa, bb in zip([rd1, sw1], [rd2, sw2], strict=True): np.testing.assert_allclose( aa.detach().cpu().numpy(), bb, diff --git a/source/tests/pt/test_utils.py b/source/tests/pt/test_utils.py index 33fe58b19b..64ea5efd21 100644 --- a/source/tests/pt/test_utils.py +++ b/source/tests/pt/test_utils.py @@ -21,7 +21,7 @@ def test_to_numpy(self) -> None: for ptp, npp in zip( [torch.float16, torch.float32, torch.float64], [np.float16, np.float32, np.float64], - strict=False, + strict=True, ): foo = foo.astype(npp) bar = to_torch_tensor(foo) From 94d2e05510ac13f1d48dededa01dd62737b063fe Mon Sep 17 00:00:00 2001 From: Copilot <198982749+Copilot@users.noreply.github.com> Date: Tue, 13 Jan 2026 00:13:37 +0800 Subject: [PATCH 12/12] fix(test): add strict parameter to zip() calls in test_adamuon.py (#221) * Initial plan * fix(test): add strict=True parameter to zip() calls in test_adamuon.py Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com> --- source/tests/pt/test_adamuon.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/source/tests/pt/test_adamuon.py b/source/tests/pt/test_adamuon.py index 4567833948..074fb67f3b 100644 --- a/source/tests/pt/test_adamuon.py +++ b/source/tests/pt/test_adamuon.py @@ -97,7 +97,9 @@ def test_optimizer_step_smoke(self) -> None: optimizer.step() # Verify all parameters with gradients changed - for i, (p, init_p) in enumerate(zip(model.parameters(), initial_params)): + for i, (p, init_p) in enumerate( + zip(model.parameters(), initial_params, strict=True) + ): if p.grad is not None: self.assertFalse( torch.allclose(p, init_p), @@ -369,7 +371,7 @@ def test_state_dict_save_load(self) -> None: params1 = list(optimizer.param_groups[0]["params"]) params2 = list(optimizer2.param_groups[0]["params"]) - for p1, p2 in zip(params1, params2): + for p1, p2 in zip(params1, params2, strict=True): s1 = optimizer.state[p1] s2 = optimizer2.state[p2] self.assertEqual(set(s1.keys()), set(s2.keys()))