fix(jax): skip padding for non-XLA SavedModels

[njzjz]

Summary

• detect whether a JAX SavedModel contains XlaCallModule during C++ initialization
• keep the dynamic atom-count padding only for XLA-compiled lower calls
• pass exact nall_real shapes for non-XLA SavedModels

This is mainly to support #5598: padding only has value for XLA static-shape execution and otherwise changes the non-XLA inference shape unnecessarily.

Tests

• git diff --check
• cmake --build source/build --target deepmd_cc -j2
• ruff check .
• ruff format --check .

Not run: JAX C++ SavedModel runtime test, because local source/tests/infer/deeppot_dpa.savedmodel is not available.

Summary by CodeRabbit

• Bug Fixes
  • Improved inference stability and correctness by applying neighbor-list input padding only for models that benefit from it.
  • Updated tensor sizing for neighbor-list computations (including coordinate/atom-type and mapping shapes) to use the appropriate atom count, reducing the risk of shape mismatches and unnecessary overhead.
  • Enhanced detection of whether the loaded computation graph uses XLA-style compilation to drive the padding behavior.

[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: 1ac15463-2f67-4fea-9ce7-1c4d5f2ab2b9

📥 Commits

Reviewing files that changed from the base of the PR and between 92f1b37 and cb22e27.

📒 Files selected for processing (2)

• source/api_cc/include/DeepPotJAX.h
• source/api_cc/src/DeepPotJAX.cc

---

📝 Walkthrough

Walkthrough

DeepPotJAX now records whether loaded TensorFlow artifacts use XLA-style compilation markers, and the neighbor-list compute path uses that flag to decide when to apply atom-count padding and matching tensor shapes.

Changes

Conditional XLA padding in DeepPotJAX

Layer / File(s)	Summary
XLA detection helpers and init flag source/api_cc/include/DeepPotJAX.h, source/api_cc/src/DeepPotJAX.cc	Adds the private XLA-usage flag, helper code that scans graph ops and serialized function bodies for XLA markers, and initialization code that stores the detection result.
Conditional padding in neighbor-list compute source/api_cc/src/DeepPotJAX.cc	Changes neighbor-list compute to derive nall_model conditionally, resize coord_double and atype to match, and pass the updated size through add_input and mapping allocation.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: conditionally skipping padding for non-XLA JAX SavedModels.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ 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

❌ Patch coverage is 63.58382% with 63 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.96%. Comparing base (a9bcbc5) to head (cb22e27).
⚠️ Report is 6 commits behind head on master.

Files with missing lines	Patch %	Lines
source/api_cc/src/DeepPotJAX.cc	63.58%	41 Missing and 22 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5602      +/-   ##
==========================================
- Coverage   82.35%   81.96%   -0.40%     
==========================================
  Files         896      959      +63     
  Lines      100952   105593    +4641     
  Branches     4059     4104      +45     
==========================================
+ Hits        83138    86547    +3409     
- Misses      16349    17559    +1210     
- Partials     1465     1487      +22     

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

[njzjz]

Moved to wanghan-iapcm's inline review thread: #5602 (comment)