fix(dpmodel): make se_e2_a type_one_side=False traceable

[stanleys12]

DescrptSeA with type_one_side=False sorted atoms by center type and sliced each type's block using bounds from int(xp.sum(atype_loc == ti)). Under make_fx with symbolic shapes (pt_expt enable_compile=True, and torch.export), that int() raises GuardOnDataDependentSymNode, so compiled training fails at the first trace.

Now I gather each center type's atoms with xp.nonzero + xp.take, run the (ti, tt) embeddings on that block, then concatenate the blocks and restore the atom order with argsort(argsort(atype_loc)). nonzero returns ascending indices, so the concatenated blocks are in stable sorted order. Each atom still goes through one embedding per neighbor type, so the cost is the same. In a quick eager CPU benchmark (nf=4, nloc=192, nnei=138, neuron [25,50,100], fwd+bwd) it runs at 1005 ms vs 1139 ms on master for ntypes=2, and 1127 ms vs 1251 ms for ntypes=8. The compressed array-API branch gets the same change. I also removed _add_to_slice since nothing calls it anymore.

jax.jit of the two-sided descriptor still fails, same as on master, because nonzero needs a static size under jit.

Tests:

• source/tests/pt_expt/test_training.py::TestTraining::test_training_loop_compiled_two_side: new. Fails on master with the error above and passes here.
• test_make_fx in source/tests/pt_expt/descriptor/test_se_e2_a.py now also covers type_one_side=False and exclude_types, including the export round trip.
• source/tests/consistent/descriptor/test_se_e2_a.py and source/tests/jax/test_model_compression.py pass locally. TF and the pt C++ op weren't installed, so those subtests were skipped. The pt_expt test_compressed_forward needs the custom op library, so only CI covers it.

Fixes #6017

Summary by CodeRabbit

• Bug Fixes
  • Improved descriptor handling when one-sided type processing is disabled, including configurations with excluded atom types or no exclusions. Atoms remain in their original order after type-based processing.
• Tests
  • Expanded checks for descriptor tracing, gradients, and export across additional configurations, and added coverage for compiled training with one-sided type processing disabled.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: d1d62fbc-76ab-4140-b31c-2cd1322e9bae

---

📥 Commits

Reviewing files that changed from the base of the PR and between 64d3266 and 8662989.

---

📒 Files selected for processing (3)

• deepmd/dpmodel/descriptor/se_e2_a.py
• source/tests/pt_expt/descriptor/test_se_e2_a.py
• source/tests/pt_expt/test_training.py

---

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

---

📝 Walkthrough

Walkthrough

The two-sided se_e2_a paths now gather center atoms by type and restore their input order without computing type boundaries from tensor values. Descriptor tracing and compiled training tests now cover type_one_side=False.

Changes

Two-sided SeA compilation

Layer / File(s)	Summary
Per-type gathering in two-sided paths deepmd/dpmodel/descriptor/se_e2_a.py	The uncompressed and compressed paths gather center atoms by type, concatenate the type blocks, and restore the original input order. The unused _add_to_slice helper was removed.
Tracing and compiled training validation source/tests/pt_expt/descriptor/test_se_e2_a.py, source/tests/pt_expt/test_training.py	The descriptor tracing test now covers both type_one_side values and excluded-type settings. A compiled training smoke test covers type_one_side=False.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix · Severity of issue fixed: Medium

Suggested reviewers: wanghan-iapcm

---

Merge Risk: ⚪ Minimal · up to 86629

No actionable merge-blocking defect was established. The change is mergeable after normal checks.

Pre-merge checks | 4 | 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 9 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely identifies the main change: making the se_e2_a type_one_side=False path traceable.
Linked Issues check		The PR satisfies the coding requirements in issue #6017. In deepmd/dpmodel/descriptor/se_e2_a.py, the type_one_side=False paths no longer convert data-dependent type counts to Python integers. The…
Out of Scope Changes check		The changes remain within issue #6017 scope. The descriptor refactor fixes the traced two-sided path, and the added make_fx and compiled-training tests verify the affected behavior. The compressed-p…

---

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🧪 Generate unit tests (beta)

• Create a new PR

---

🛠️ Fix failing CI checks 💡

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

---

_{Comment @coderabbitai help to get the list of available commands.}