fix(dpmodel): move econf embeddings to active backend#5840
Conversation
Materialize portable NumPy electronic-configuration tables in the namespace, dtype, and device of the type-embedding weights before evaluation. This prevents NativeLayer from selecting NumPy and coercing backend-native parameters back to CPU. Add a direct mixed-state regression covering Torch output identity, numerical parity, and change_type_map regeneration without relying on eager-converting backend wrappers. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
Warning Review limit reached
Next review available in: 39 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthrough
ChangesElectronic-configuration backend alignment
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5840 +/- ##
==========================================
+ Coverage 78.58% 78.63% +0.04%
==========================================
Files 1050 1054 +4
Lines 120637 121758 +1121
Branches 4356 4406 +50
==========================================
+ Hits 94801 95741 +940
- Misses 24278 24444 +166
- Partials 1558 1573 +15 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The fix is correct and minimal — materializing econf_tebd onto the active backend at call time resolves the namespace-coercion, preserves the numpy storage for serialization, and (verified locally) also repairs the JAX nnx.jit/nnx.grad crash that use_econf_tebd=True hit before. One coverage note inline.
Exercise both one-hot and electronic-configuration type embeddings under nnx.jit and nnx.grad, including the type map required by the electronic path. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@source/tests/jax/test_type_embed.py`:
- Around line 41-44: Strengthen the gradient check in the test around forward
and grad by capturing the result of grad(type_embedding) and validating its
gradient leaves are finite and non-degenerate where applicable, rather than only
asserting the gradient state is non-null. Preserve the existing output-shape and
output-NaN checks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 823f87e8-3fcc-4a1b-b2ff-0d54a1deb717
📒 Files selected for processing (1)
source/tests/jax/test_type_embed.py
Assert that traced TypeEmbedNet gradient states contain finite, nonzero leaves for both one-hot and electronic-configuration embeddings. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
This closes the gap properly. Adding use_econf_tebd=True to the nnx.jit/nnx.grad test covers the real production path rather than a constructed one: before this, the only new test was torch-only and eager, and it manufactured a numpy-econf-beside-torch-weights state that no wrapper actually produces (pt_expt stores econf_tebd as a torch buffer, JAX as an ArrayAPIVariable), so it guarded the class rather than the failure.
That matters for whether this PR has a genuine regression test, and now it does. I reproduced the pre-fix crash locally - use_econf_tebd=True under nnx.jit threw DynamicJaxprTracer has no attribute device, the same pattern #5736 fixed for the sibling xp.eye branch - so the new subtest fails before the fix and passes after. The any(np.abs(leaf) > 0.0) assertion added in the follow-up is a good touch; it stops the test passing on a degenerate all-zero gradient.
One optional follow-up, not for this PR: change_type_map() still calls array_api_compat.device() directly instead of the tracing-safe _array_device_or_none() helper, which is now the last instance of that pattern in the file.
3 participants
Fixes #5641
Summary
NativeLayerfrom selecting NumPy and coercing Torch/JAX/array-API weights back to CPUchange_type_map()through the same call-time boundaryWhy existing tests missed this
The cross-backend consistency tests already include electronic-configuration embeddings, but their JAX, array-api-strict, TF2, and exportable-PyTorch wrappers eagerly convert NumPy attributes when constructing the backend object. That converts both the weights and
econf_tebd, so the generic dpmodel mixed-state boundary is never exercised.The direct dpmodel path uses NumPy parameters and a NumPy electronic-configuration table, where
NativeLayerselecting its namespace from the input is correct. Existing tests therefore compared homogeneous NumPy state or homogeneous wrapper-converted state, but never backend-native weights alongside the intentionally portable NumPy constant. CPU-only numerical comparisons also did not assert the returned array backend, dtype, or device.Validation
source/tests/consistent/test_type_embedding.py: 64 passed, 49 backend-dependent skipschange_type_map()regenerates the NumPy electronic-configuration tableruff format .ruff check .Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
Bug Fixes
Tests