fix(jax): return stored HLO neighbor count

[njzjz]

Summary

• make HLO.get_nnei() derive the total neighbor width from the stored sel metadata;
• delegate to get_nsel() so the metadata contract has one source of truth;
• add a focused regression for a serialized HLO model's selection metadata.

Why existing tests missed this

The constructor stores self.sel but never creates self.nsel. Existing consistency and model tests exercise neighbor counts on dpmodel, TensorFlow, and PyTorch wrappers, but no JAX test directly queried the HLO wrapper. Serialization and export helpers can therefore reach an AttributeError even though the broader model tests pass.

The new unit test isolates the metadata API without constructing real StableHLO byte artifacts, which are unrelated to this accessor.

Validation

• targeted HLO metadata test: 1 passed;
• pre-fix implementation fails the same test with AttributeError;
• ruff format .;
• ruff check .;
• git diff --check.

Closes #5669.

Coding agent: Codex
Codex version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning effort: xhigh

Summary by CodeRabbit

• Bug Fixes
  • Corrected neighbor-count reporting for serialized JAX HLO model metadata.
  • Neighbor counts now accurately reflect the stored selection values.
• Tests
  • Added regression coverage to verify neighbor-count calculations.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: d3e0561e-37b3-4560-b729-57a8b630f4a2

📥 Commits

Reviewing files that changed from the base of the PR and between eacab7c and caf5523.

📒 Files selected for processing (1)

• source/tests/jax/test_hlo.py

🚧 Files skipped from review as they are similar to previous changes (1)

• source/tests/jax/test_hlo.py

---

📝 Walkthrough

Walkthrough

HLO.get_nnei() now delegates to get_nsel(), and a regression test verifies the neighbor count derived from stored sel metadata.

Changes

JAX HLO neighbor count

Layer / File(s)	Summary
Neighbor count accessor and regression coverage deepmd/jax/model/hlo.py, source/tests/jax/test_hlo.py	HLO.get_nnei() uses get_nsel(), and a regression test asserts the result equals the sum of sel.

Estimated code review effort: 1 (Trivial) | ~5 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes the JAX HLO neighbor-count fix and matches the main change.
Linked Issues check	✅ Passed	The PR updates HLO.get_nnei() to use stored selection metadata and adds a regression test, matching issue #5669.
Out of Scope Changes check	✅ Passed	The changes stay focused on the HLO neighbor-count fix and its test, with no unrelated edits.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.57%. Comparing base (7d5ad38) to head (caf5523).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5792      +/-   ##
==========================================
- Coverage   79.85%   79.57%   -0.28%     
==========================================
  Files        1022     1022              
  Lines      117351   117350       -1     
  Branches     4313     4312       -1     
==========================================
- Hits        93706    93386     -320     
- Misses      22101    22407     +306     
- Partials     1544     1557      +13     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.