mxtensor: switch to AOBaseTensor dispatch #3080
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
|
Stack from ghstack (oldest at bottom): |
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3080
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 3467ee3 with merge base b3b545f ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
This was referenced Sep 26, 2025
vkuzo
added a commit
that referenced
this pull request
Sep 26, 2025
Summary: Deletes the custom op dispatch logic in `MXTensor` and switches to the general one in `AOBaseTensor`. To enable this, we move the mx ops to the same file as `MXTensor`. This is to avoid the need for callsites to separately import `mx_ops.py`. Note that there are a couple of custom mx ops which could instead use the general implementation, leaving that for future PRs to keep this one small. Test Plan: ``` pytest test/prototype/mx_formats -s -x ``` Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 1d27ad4 ghstack-comment-id: 3338749802 Pull Request resolved: #3080
vkuzo
added a commit
that referenced
this pull request
Sep 26, 2025
Summary: Deletes the custom op dispatch logic in `MXTensor` and switches to the general one in `AOBaseTensor`. To enable this, we move the mx ops to the same file as `MXTensor`. This is to avoid the need for callsites to separately import `mx_ops.py`. Note that there are a couple of custom mx ops which could instead use the general implementation, leaving that for future PRs to keep this one small. Test Plan: ``` pytest test/prototype/mx_formats -s -x ``` Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 994257e ghstack-comment-id: 3338749802 Pull Request resolved: #3080
vkuzo
added
the
topic: improvement
jerryzh168
reviewed
Sep 26, 2025
| implements = MXTensor.implements | ||
|
|
||
|
|
||
| @implements([aten.detach.default, aten.alias.default]) |
There was a problem hiding this comment.
these are defined in TorchAOBaseTensor as well I think
There was a problem hiding this comment.
yes, we can move in future PRs if needed, I wanted to minimize changes in this PR
jerryzh168
approved these changes
Sep 26, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 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.
Summary:
Deletes the custom op dispatch logic in
MXTensorand switches to thegeneral one in
AOBaseTensor.To enable this, we move the mx ops to the same file as
MXTensor. Thisis to avoid the need for callsites to separately import
mx_ops.py.Note that there are a couple of custom mx ops which could instead use
the general implementation, leaving that for future PRs to keep this one
small.
Test Plan:
Reviewers:
Subscribers:
Tasks:
Tags: