[AINode] Replace DiT-derived Sundial conditioning head - #18616
Merged
CRZbulabula merged 1 commit intoSep 10, 2026
Conversation
Contributor
There was a problem hiding this comment.
🟢 Approval recommended
The replacement preserves the existing computation and checkpoint contract with focused automated coverage.
Pull request overview
Replaces DiT-derived Sundial conditioning components with independently implemented, checkpoint-compatible PyTorch modules.
Changes:
- Reimplements conditional residual and output projection layers.
- Preserves parameter names, initialization, shapes, and behavior.
- Adds focused compatibility and behavior tests.
File summaries
| File | Description |
|---|---|
iotdb-core/ainode/iotdb/ainode/core/model/sundial/flow_loss.py |
Replaces DiT-derived conditioning layers. |
iotdb-core/ainode/tests/test_sundial_flow_loss.py |
Tests checkpoint compatibility, gradients, initialization, and sampling. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 0
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
HTHou
approved these changes
Sep 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR resolves the Sundial provenance issue identified in the review of #18611. The prior conditioning head was explicitly identified as adopted from DiT, whose CC-BY-NC-4.0 license is not suitable for Apache source distributions.
The DiT-derived residual and output-layer implementation is replaced with an independently written PyTorch implementation based on Sundial's required tensor and checkpoint contract. The replacement:
No user-facing behavior change is intended. Class names and computation are expressed independently, while nested module names remain stable solely for checkpoint compatibility.
Focused tests cover the checkpoint key/shape contract, strict state-dict loading, output shapes, gradient propagation, zero-initialized output, and sampling shape. A random-weight comparison against the previous implementation also produced identical output with a maximum absolute delta of
0.0.Verification
PYTHONPATH=. python3 -m unittest discover -s tests -p "test_*.py" -v(7 tests passed)black --check .with Black 25.1.0isort --check-only --profile black .with isort 6.0.1python3 -m compileall -q iotdb-core/ainode/iotdb iotdb-core/ainode/testsmvn -q -P with-ainode apache-rat:check -pl iotdb-core/ainode -DskipTestsThis PR has:
Key changed/added classes (or packages if there are too many classes) in this PR
iotdb.ainode.core.model.sundial.flow_losstests.test_sundial_flow_loss