fix(jax): setattr case_embd

[njzjz]

case_embd was supported but the JAX backend was not touched.

Summary by CodeRabbit

• Bug Fixes

  • Ensured the case_embd parameter is consistently converted and handled during fitting, improving compatibility across array backends and preventing mis-coercion.

• Tests

  • Adjusted test setup to reset the default computation graph before enabling eager execution, stabilizing related test runs.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

📝 Walkthrough

Walkthrough

Added "case_embd" to fitting attribute conversion logic for JAX and strict Array API paths, and updated a TensorFlow test module to reset the default graph before enabling eager execution.

Changes

Cohort / File(s)	Change Summary
JAX fitting update deepmd/jax/fitting/fitting.py	Added "case_embd" to the set of parameter names routed through to_jax_array and wrapped/converted via the existing general-fitting setter logic.
Array-API strict fitting tests source/tests/array_api_strict/fitting/fitting.py	Added "case_embd" to the set of names passed to to_array_api_strict_array within the general fitting attribute setter in tests.
TensorFlow test setup source/tests/pt/test_tabulate.py	Inserted tf.reset_default_graph() at the start of setUpModule() before enabling eager execution.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

• Check that case_embd is a valid fitting parameter for the affected modules.
• Verify conversions behave correctly for None and non-None values across supported Flax/array-API versions.
• Confirm the TF reset call is safe in the test environment and doesn't interfere with other tests.

Pre-merge checks and finishing touches

❌ 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%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title directly relates to the main change: adding case_embd support to the JAX backend's setattr logic.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 31f8647 and d4277dd.

📒 Files selected for processing (1)

• source/tests/pt/test_tabulate.py (1 hunks)

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

• source/tests/pt/test_tabulate.py

---

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 and usage tips.}

[Copilot]

Pull request overview

This PR fixes a missing attribute handler for case_embd in the JAX backend. The case_embd attribute was already supported in PyTorch and Paddle backends but was not being properly handled in the JAX backend's setattr_for_general_fitting function.

• Added case_embd to the set of attributes that are converted to JAX arrays in the fitting module

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.14%. Comparing base (26013cb) to head (d4277dd).
⚠️ Report is 4 commits behind head on devel.

Additional details and impacted files

@@           Coverage Diff           @@
##            devel    #5104   +/-   ##
=======================================
  Coverage   82.14%   82.14%           
=======================================
  Files         709      709           
  Lines       72458    72468   +10     
  Branches     3615     3615           
=======================================
+ Hits        59520    59530   +10     
+ Misses      11776    11775    -1     
- Partials     1162     1163    +1     

☔ View full report in Codecov by Sentry.
📢 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.