Use a different str for device instead of meta #7995
Merged
Conversation
This file contains 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
qihqi
force-pushed
the
hanq_xla2_2
branch
2 times, most recently
from
d8fb3de
to
e099f54
Compare
So XLATensor2's device property is no longer 'meta' Capture tensor constructors and give option to override 'cuda' with 'jax' device (specially useful for running GPU code on TPU). Co-authored-by: Will Cromar <[email protected]>
ManfeiBai
reviewed
Sep 12, 2024
| @@ -1634,6 +1645,7 @@ def _aten_arange( | |||
| requires_grad=False, | |||
| device=None, | |||
| pin_memory=False, | |||
| **kwargs | |||
| ): | |||
There was a problem hiding this comment.
do we need add **kwargs here to accept **kwargs from res2 = func(input2, *args2, **kwargs2)?
I seems met one situation that: func in jaten.py file didn't received kwargs2 from res2 = func(input2, *args2, **kwargs2)
There was a problem hiding this comment.
yeah this is not needed. removed.
ManfeiBai
approved these changes
Sep 12, 2024
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.
Implements this idea: #7705 but without the C++ piece.
We could use anything for the device (incl just a string: https://github.com/albanD/subclass_zoo/blob/main/new_device.py); however it's advantagous to use a torch.device with an existing registered key. (we could use 'xla:0' too).
The caveat is that we need to capture tensor constructors in jtorch instead of jaten (like it used to be) because otherwise torch will check if such device module is loaded or not (example for capturing 'cuda').
fixes #7966