Separated implementation of aten::scatter overloads #2605
Merged
Conversation
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
… for the value argument in aten::scatter.value Signed-off-by: Linsho Kaku <[email protected]>
|
@microsoft-github-policy-service agree |
xadupre
approved these changes
Oct 6, 2025
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2605 +/- ##
==========================================
+ Coverage 69.94% 69.95% +0.01%
==========================================
Files 222 222
Lines 26307 26311 +4
Branches 2604 2604
==========================================
+ Hits 18400 18406 +6
+ Misses 6995 6993 -2
Partials 912 912 ☔ View full report in Codecov by Sentry. |
justinchuby
added
the
module: torchlib
|
Thanks for the high quality contribution! |
justinchuby
reviewed
Oct 6, 2025
| def aten_scatter( | ||
| @torch_op("aten::scatter.src", trace_only=True) | ||
| def aten_scatter_src( | ||
| self: TReal, |
There was a problem hiding this comment.
One thing I just noticed:
Suggested change
| self: TReal, | |
| self: TensorType, |
We should update the accepted type to make it more permissive. Could you create a follow PR?
justinchuby
reviewed
Oct 6, 2025
| return op.ScatterElements(self, index, update, axis=dim) | ||
| @torch_op("aten::scatter.value", trace_only=True) | ||
| def aten_scatter_value( | ||
| self: TReal, |
There was a problem hiding this comment.
Suggested change
| self: TReal, | |
| self: TensorType, |
justinchuby
reviewed
Oct 6, 2025
| self: TReal, | ||
| dim: int, # we have to use int here because ScatterElements() will use this attribute | ||
| index: TInt, | ||
| value: TReal, |
There was a problem hiding this comment.
Suggested change
| value: TReal, | |
| value: float, |
justinchuby
reviewed
Oct 6, 2025
Comment on lines
+7759
to
+7760
| scalar_tensor = op.CastLike(value, self) | ||
| src = op.Expand(scalar_tensor, op.Shape(index)) |
There was a problem hiding this comment.
This can be simplified:
Suggested change
| scalar_tensor = op.CastLike(value, self) | |
| src = op.Expand(scalar_tensor, op.Shape(index)) | |
| scalar_tensor = ir.tensor([value], dtype=self.dtype) | |
| src = op.ConstantOfShape(op.Shape(index), value=scalar_tensor) |
|
@linshokaku I added some suggestions after the PR was merged. Could you create a follow PR to address them? Thanks a lot. |
justinchuby
pushed a commit
that referenced
this pull request
Oct 7, 2025
…lue (#2612) follow #2605 --------- Signed-off-by: Linsho Kaku <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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.
close #2601 #2602
This PR refactors the implementation of
aten::scatteroverloads, improving the clarity of the ONNX output generated byaten::scatter.src.I've also added new tests to verify the correctness of these changes. To make the added tests pass, I needed to also address the issue reported in #2602, which is included in this PR's diff.