[Trace Python dispatcher] Support torch.DispatchKey & torch.DispatchKeySet #144439
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/144439
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 2 Unrelated FailuresAs of commit 8732618 with merge base d44c390 ( NEW FAILURE - The following job has failed:
FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
torch/_dynamo/variables/torch.py
Outdated
| @classmethod | ||
| def create_with_source(cls, value, source): | ||
| install_guard(source.make_guard(GuardBuilder.ID_MATCH)) | ||
| return cls(value, source=source) |
There was a problem hiding this comment.
DispatchKeys with the same items in them do not end up with the same ID_MATCH because they can be different objects. Is there something like an INT_MATCH we could do on the integer representation of the DispatchKeySet?
There was a problem hiding this comment.
Right, I think there are two options here:
- Add a
.python_setattribute toDispatchKeySetand guard on it. However, I’m uncertain about its potential impact on existing eager mode logic. - Introduce a new guard type specifically for
DispatchKeySetin Dynamo. But I'm not sure if this is overengineering that we have to implement guard for every new custom data type.
There was a problem hiding this comment.
Just chatted with @anijain2305, we already checked DispatchKeySet in the TENSOR_MATCH guard:
pytorch/torch/csrc/dynamo/guards.cpp
Line 113 in 6bc17b0
So probably I can add the
DispatchKeySet guard by leveraging that.
Pull Request resolved: #144444 Approved by: https://github.com/williamwen42, https://github.com/zou3519 ghstack dependencies: #144439
Stack from ghstack (oldest at bottom):
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames