Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove non-parallel approaches to get a device #8827

Open
yaoshiang opened this issue Mar 12, 2025 · 1 comment
Open

Remove non-parallel approaches to get a device #8827

yaoshiang opened this issue Mar 12, 2025 · 1 comment
Labels
enhancement New feature or request usability Bugs/features related to improving the usability of PyTorch/XLA

Comments

@yaoshiang
Copy link

yaoshiang commented Mar 12, 2025

🐛 Bug

There are four ways to get a device object. Three could be deprecated as they are likely unnecessary, and also anti-parallel to torch.cuda.device(), which does not return a device but is in fact a context manager.

torch.device('xla', 3)
torch_xla.device(3)
torch_xla.torch_xla.device(3)
torch_xla.core.xla_model.xla_device(3)

To Reproduce

Run this in a collab with a TPU enabled and pip installed pytorch_xla with the correct version.

torch.device('xla', 3)
torch_xla.device(3)
torch_xla.torch_xla.device(3)
torch_xla.core.xla_model.xla_device(3)

Expected behavior

Should deprecate the three ways to get a device inside torch_xla.
Docs that use them should be updated to torch.device("xla", <optional int>) and remove import of torch_xla.core.xla_model if necessary.

Environment

TPU on Collab

Additional context

Should verify that the four APIs do ultimately call the same underlying function to guarantee equivalent behavior.

@ysiraichi ysiraichi added enhancement New feature or request usability Bugs/features related to improving the usability of PyTorch/XLA labels Mar 14, 2025
@ysiraichi
Copy link
Collaborator

I agree. We should not have this many ways to get an XLA device. That said, I'm not against having one internal torch_xla.device(3) function, so that we have control over that.

cc @miladm @zpcore @tengyifei @ManfeiBai

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request usability Bugs/features related to improving the usability of PyTorch/XLA
Projects
None yet
Development

No branches or pull requests

2 participants