fix(dpmodel): materialize fitting buffers on active backend#5841
Conversation
Convert portable fitting constants locally into the runtime namespace, dtype, and device before default-fparam handling, parameter normalization, case embedding, atomic bias gathering, and polar scale/shift postprocessing. Add direct dpmodel regressions with Torch runtime inputs while all nine persistent buffers remain NumPy, verifying backend identity and numerical parity without eager wrapper conversion. 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: 36 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)
📝 WalkthroughWalkthroughGeneralFitting and PolarFitting now materialize stored buffers on the active backend, dtype, and device before runtime operations. New Torch and JAX tests validate backend outputs and tracing behavior. ChangesBackend buffer materialization
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: 🚥 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 #5841 +/- ##
==========================================
+ Coverage 78.58% 78.63% +0.05%
==========================================
Files 1050 1054 +4
Lines 120637 121768 +1131
Branches 4356 4404 +48
==========================================
+ Hits 94801 95754 +953
- Misses 24278 24441 +163
- 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 — all 9 buffers are materialized to the right float runtime reference's dtype/device, self.<buffer> is never mutated (numpy serialization preserved), and I confirmed the JAX side is safe: array_api_compat.device() on a traced tensor returns None (not a crash), so the pattern runs cleanly under jax.jit/jax.grad. One coverage note inline.
Exercise default frame parameters, frame/atom normalization buffers, case embeddings, and atom biases through an EnergyFittingNet forward and gradient under nnx.jit. 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_fitting.py`:
- Around line 30-37: Strengthen the fitting test around the forward/gradient
invocation by snapshotting every configured portable buffer, including values
and namespace/type, then assert they are unchanged after execution. Add a
reference or differential-output check demonstrating that changing each
non-default buffer affects the result, while preserving the existing shape, NaN,
and gradient assertions.
🪄 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: 70c6a27b-d4c3-4b96-8246-9e0af1df9dbb
📒 Files selected for processing (1)
source/tests/jax/test_fitting.py
Assert portable fitting buffers retain identity, namespace, and values across traced forward/gradient calls. Neutralize each configured buffer independently to prove it affects the output. 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.
The JAX-tracing gap is closed: test_runtime_buffers_support_nnx_jit_and_grad_tracing exercises the forward under nnx.jit and takes a parameter gradient through nnx.grad, with non-default values for every materialized buffer and fparam omitted so default_fparam_tensor is covered too. Both commits are test-only, no production change.
The follow-up test is a better one than I asked for. Asserting assertIs(current, original) on each buffer pins the numpy-storage-preserved invariant, and substituting a neutral value per buffer and requiring the output to change proves each materialization is load-bearing rather than the test passing vacuously.
For the record on the earlier concern about this reintroducing the #5736 device-lookup crash: I checked it directly rather than by inspection - array_api_compat.device() returns None on a JAX tracer, and this PR's xp.asarray(buf, dtype=..., device=device(t)) pattern runs clean under both jax.jit and jax.grad. The difference from #5736 is that the device is taken from a computed runtime value rather than a raw NNX variable. The new test locks that behaviour in.
c63bb90
3 participants
Fixes #5642
Summary
PolarFitting.scaleandconstant_matrix, which had the identical failure mode during polar postprocessingThe complete affected set is nine arrays:
default_fparam_tensor,fparam_avg,fparam_inv_std,aparam_avg,aparam_inv_std,case_embd,bias_atom_e, polarscale, and polarconstant_matrix.Why existing tests missed this
Cross-backend fitting tests already exercise default fparam, parameter statistics, biases, and polar outputs, but they instantiate backend wrappers that eagerly convert every NumPy attribute. Those objects are homogeneous before forward execution and therefore never expose the generic dpmodel boundary.
Direct dpmodel fitting tests use NumPy descriptors, atom types, parameters, and buffers, so the selected namespace is also NumPy and all operations are valid. Dedicated PyTorch, Paddle, and TensorFlow implementations maintain their own backend-native state and do not exercise this shared dpmodel implementation. No previous test combined a raw dpmodel fitting object, backend-native runtime inputs, and still-NumPy portable buffers.
Validation
ruff 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