From 866298955aba6e773022f458db2fc1b47df68af9 Mon Sep 17 00:00:00 2001 From: Stanley Shen Date: Mon, 12 Oct 2026 11:26:14 -0700 Subject: [PATCH] fix(dpmodel): make se_e2_a type_one_side=False traceable The two-sided path sliced per-type blocks with bounds taken from int(xp.sum(atype_loc == ti)), which cannot be traced with symbolic shapes, so pt_expt compiled training and export failed. Gather each center type's atoms with nonzero instead and restore the order with the inverse of the stable sort. The embedding still runs once per atom. Co-Authored-By: Claude Opus 5.5 --- deepmd/dpmodel/descriptor/se_e2_a.py | 95 ++++++------------- .../tests/pt_expt/descriptor/test_se_e2_a.py | 6 +- source/tests/pt_expt/test_training.py | 10 ++ 3 files changed, 46 insertions(+), 65 deletions(-) diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index 521dc3246c..f22af56b51 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -477,14 +477,6 @@ def _tabulate_fusion_se_a( ) return xp.sum(em[:, :, :, None] * values[:, :, None, :], axis=1) - def _add_to_slice(self, array: Array, start: int, end: int, value: Array) -> Array: - """Return ``array`` with ``value`` added to ``array[start:end]``.""" - xp = array_api_compat.array_namespace(array, value) - return xp.concat( - [array[:start], array[start:end] + value, array[end:]], - axis=0, - ) - def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], @@ -578,41 +570,28 @@ def call( gg = self.cal_g(ss, (tt,)) gr += xp.sum(gg[:, :, :, None] * tr[:, :, None, :], axis=1) else: - # Sort atoms by center type so each type forms a contiguous block. - # Slice indexing (arr[s:e]) is array-api compatible and lets us - # run cal_g only on atoms of the matching center type, keeping the - # same O(nf*nloc) total embedding cost as the original numpy code. + # Gather the atoms of each center type so cal_g only runs on them. + # nonzero keeps the block sizes traceable, unlike int() on a sum. + # Its indices are ascending, so the concatenated blocks are in + # (stable) sorted order. atype_loc = xp.reshape(atype_ext[:, :nloc], (nf * nloc,)) - sort_idx = xp.argsort(atype_loc) - unsort_idx = xp.argsort(sort_idx) - rr_s = xp.take(rr, sort_idx, axis=0) - mask_s = xp.take(exclude_mask, sort_idx, axis=0) + unsort_idx = xp.argsort(xp.argsort(atype_loc)) dev = array_api_compat.device(coord_ext) - gr_s = xp.zeros([nf * nloc, ng, 4], dtype=input_dtype, device=dev) - # Per-type boundaries in sorted order - type_ends = [] - offset = 0 + gr_blocks = [] for ti in range(self.ntypes): - offset += int(xp.sum(xp.astype(atype_loc == ti, xp.int32))) - type_ends.append(offset) - type_starts = [0, *type_ends[:-1]] - for ti in range(self.ntypes): - s, e = type_starts[ti], type_ends[ti] - if s == e: - continue + idx = xp.nonzero(atype_loc == ti)[0] + rr_t = xp.take(rr, idx, axis=0) + mask_t = xp.take(exclude_mask, idx, axis=0) + gr_t = xp.zeros([idx.shape[0], ng, 4], dtype=input_dtype, device=dev) for tt in range(self.ntypes): - mm = mask_s[s:e, sec[tt] : sec[tt + 1]] - tr = rr_s[s:e, sec[tt] : sec[tt + 1], :] + mm = mask_t[:, sec[tt] : sec[tt + 1]] + tr = rr_t[:, sec[tt] : sec[tt + 1], :] tr = tr * xp.astype(mm[:, :, None], tr.dtype) ss = tr[..., 0:1] gg = self.cal_g(ss, (ti, tt)) - gr_s = self._add_to_slice( - gr_s, - s, - e, - xp.sum(gg[:, :, :, None] * tr[:, :, None, :], axis=1), - ) - gr = xp.take(gr_s, unsort_idx, axis=0) + gr_t += xp.sum(gg[:, :, :, None] * tr[:, :, None, :], axis=1) + gr_blocks.append(gr_t) + gr = xp.take(xp.concat(gr_blocks, axis=0), unsort_idx, axis=0) gr = xp.reshape(gr, (nf, nloc, ng, 4)) # nf x nloc x ng x 4 gr /= self.nnei @@ -664,45 +643,33 @@ def _call_compressed( ) else: atype_loc = xp.reshape(atype_ext[:, :nloc], (nfnl,)) - sort_idx = xp.argsort(atype_loc) - unsort_idx = xp.argsort(sort_idx) - rr_s = xp.take(rr, sort_idx, axis=0) - mask_s = xp.take(exclude_mask, sort_idx, axis=0) + unsort_idx = xp.argsort(xp.argsort(atype_loc)) dev = array_api_compat.device(rr) - gr_s = xp.zeros([nfnl, 4, ng], dtype=rr.dtype, device=dev) - type_ends = [] - offset = 0 - for ti in range(self.ntypes): - offset += int(xp.sum(xp.astype(atype_loc == ti, xp.int32))) - type_ends.append(offset) - type_starts = [0, *type_ends[:-1]] + gr_blocks = [] for ti in range(self.ntypes): - s, e = type_starts[ti], type_ends[ti] - if s == e: - continue + idx = xp.nonzero(atype_loc == ti)[0] + rr_t = xp.take(rr, idx, axis=0) + mask_t = xp.take(exclude_mask, idx, axis=0) + gr_t = xp.zeros([idx.shape[0], 4, ng], dtype=rr.dtype, device=dev) for tt in range(self.ntypes): embedding_idx = tt * self.ntypes + ti compress_data_ii = self.compress_data[embedding_idx] if array_api_compat.size(compress_data_ii) == 0: continue compress_info_ii = self.compress_info[embedding_idx] - mm = mask_s[s:e, sec[tt] : sec[tt + 1]] - rr_i = rr_s[s:e, sec[tt] : sec[tt + 1], :] + mm = mask_t[:, sec[tt] : sec[tt + 1]] + rr_i = rr_t[:, sec[tt] : sec[tt + 1], :] rr_i = rr_i * xp.astype(mm[:, :, None], rr_i.dtype) ss = rr_i[:, :, :1] - gr_s = self._add_to_slice( - gr_s, - s, - e, - self._tabulate_fusion_se_a( - compress_data_ii, - compress_info_ii, - ss, - rr_i, - ng, - ), + gr_t += self._tabulate_fusion_se_a( + compress_data_ii, + compress_info_ii, + ss, + rr_i, + ng, ) - gr = xp.take(gr_s, unsort_idx, axis=0) + gr_blocks.append(gr_t) + gr = xp.take(xp.concat(gr_blocks, axis=0), unsort_idx, axis=0) return xp.permute_dims(gr, (0, 2, 1)) def serialize(self) -> dict: diff --git a/source/tests/pt_expt/descriptor/test_se_e2_a.py b/source/tests/pt_expt/descriptor/test_se_e2_a.py index e3a8ca5c21..3244b28f23 100644 --- a/source/tests/pt_expt/descriptor/test_se_e2_a.py +++ b/source/tests/pt_expt/descriptor/test_se_e2_a.py @@ -164,8 +164,10 @@ def test_compressed_forward(self, tos, prec) -> None: assert result_ref.shape == result_cmp.shape torch.testing.assert_close(result_ref, result_cmp, atol=atol, rtol=atol) + @pytest.mark.parametrize("em", [[], [[0, 1]]]) # exclude_types + @pytest.mark.parametrize("tos", [True, False]) # type_one_side @pytest.mark.parametrize("prec", ["float64"]) # precision - def test_make_fx(self, prec) -> None: + def test_make_fx(self, prec, tos, em) -> None: rng = np.random.default_rng(GLOBAL_SEED) _, _, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) @@ -179,6 +181,8 @@ def test_make_fx(self, prec) -> None: self.rcut_smth, self.sel, precision=prec, + type_one_side=tos, + exclude_types=em, seed=GLOBAL_SEED, ).to(self.device) dd0.davg = torch.tensor(davg, dtype=dtype, device=self.device) diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index 4fdad73dc8..5863f0a38b 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -629,6 +629,16 @@ def test_training_loop_compiled_silu(self) -> None: config = normalize(config) self._run_training(config) + @REQUIRES_SUPPORTED_COMPILE + def test_training_loop_compiled_two_side(self) -> None: + """Run compiled training with type_one_side=False.""" + config = _make_config(self.data_dir, numb_steps=5) + config["model"]["descriptor"]["type_one_side"] = False + config["training"]["enable_compile"] = True + config = update_deepmd_input(config, warning=False) + config = normalize(config) + self._run_training(config) + class TestCompiledModelGetattr(unittest.TestCase): """Unit tests for _CompiledModel attribute delegation.