8000 Fix `MaskedTensor` to device ignored mask by zeshengzong · Pull Request #151205 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

Fix MaskedTensor to device ignored mask #151205

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.

Sign up for GitHub

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

Jump to bottom
Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Fix MaskedTensor to device ignored mask #151205

wants to merge 3 commits into from

Conversation

zeshengzong
Copy link
Contributor

Fixes #147140

Changes

  • Add to implementation in MaskedTensor to support move mask to target device

Test Result

In [1]: import torch
   ...: from torch.masked import as_masked_tensor
   ...: data = torch.tensor([1,2,3])
   ...: mask = torch.tensor([True,False,True])
   ...: mt = as_masked_tensor(data<
8000
/span>, mask).to('cuda')
   ...: mt.get_data().device, mt.get_mask().device
/home/zong/code/pytorch/torch/masked/maskedtensor/core.py:247: UserWarning: The PyTorch API of MaskedTensors is in prototype stage and will change in the near future. Please open a Github issue for features requests and see our documentation on the torch.masked module for further information about the project.
  return MaskedTensor(data, mask)
/home/zong/code/pytorch/torch/masked/maskedtensor/_ops_refs.py:354: UserWarning: The PyTorch API of MaskedTensors is in prototype stage and will change in the near future. Please open a Github issue for features requests and see our documentation on the torch.masked module for further information about the project.
  return MaskedTensor(new_data, _maybe_get_mask(args[0]))
Out[1]: (device(type='cuda', index=0), device(type='cuda', index=0))

In [2]: mt.sum(dim=0)
/home/zong/code/pytorch/torch/masked/maskedtensor/core.py:247: UserWarning: The PyTorch API of MaskedTensors is in prototype stage and will change in the near future. Please open a Github issue for features requests and see our documentation on the torch.masked module for further information about the project.
  return MaskedTensor(data, mask)
Out[2]: MaskedTensor(4, True)
pytest test/test_maskedtensor.py -vv

image

@pytorch-bot PyTorch Bot
Copy link
pytorch-bot bot commented Apr 14, 2025 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/151205

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit b81d648 with merge base 25803d3 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@zeshengzong zeshengzong marked this pull request as ready for review April 14, 2025 07:37
@zeshengzong
Copy link
Contributor Author

@pytorchbot label "topic: not user facing"

@pytorch-bot pytorch-bot bot added the topic: not user facing topic category label Apr 14, 2025
Skylion007
Skylion007 reviewed Apr 14, 2025
View reviewed changes
torch/masked/maskedtensor/core.py Outdated
device_to = current_device
if len(args) == 1:
arg = args[0]
if isinstance(arg, (torch.device, str, int)):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use DeviceLIke from torch._prims_common and maybe use cannonicalize_device below?

Copy link
Contributor Author
@zeshengzong zeshengzong Apr 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi, do you mean using DeviceLikeType like this?

    def to(self, *args, **kwargs):
        current_device = self._masked_data.device
        device_to = current_device
        if len(args) == 1:
            arg = args[0]
            if isinstance(arg, DeviceLikeType):
                device_to = canonicalize_device(arg)

But mypy seems unhapply about it

>>> Lint for torch/masked/maskedtensor/core.py:

  Error (MYPY) [misc]
    Parameterized generics cannot be used with class or instance checks

        364  |        device_to = current_device
        365  |        if len(args) == 1:
        366  |            arg = args[0]
    >>> 367  |            if isinstance(arg, DeviceLikeType):
        368  |                device_to = canonicalize_device(arg)
        369  |            elif isinstance(arg, torch.Tensor):
        370  |                device_to = arg.device

  Error (MYPY) [arg-type]
    Argument 2 to "isinstance" has incompatible type "<typing special form>";
    expected "_ClassInfo"

        364  |        device_to = current_device
        365  |        if len(args) == 1:
        366  |            arg = args[0]
    >>> 367  |            if isinstance(arg, DeviceLikeType):
        368  |                device_to = canonicalize_device(arg)
        369  |            elif isinstance(arg, torch.Tensor):
        370  |                device_to = arg.device

And found a solution in here python/mypy#12155 (comment)

@zeshengzong zeshengzong force-pushed the fix/maskedtensor/to branch from 15fce79 to eea4367 Compare April 15, 2025 01:53
@zeshengzong zeshengzong changed the title Fix MaskedTensor to device ignored mask` Fix MaskedTensor to device ignored mask Apr 15, 2025
@mikaylagawarecki mikaylagawarecki added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Apr 16, 2025
Skylion007
8000 Skylion007 reviewed May 12, 2025
View reviewed changes
torch/masked/maskedtensor/core.py
device_to = current_device
if len(args) == 1:
arg = args[0]
if isinstance(arg, get_args(DeviceLikeType)):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This actually makes typing worse unfortunately as it breaks the type narrower so what you had before was probably better.

@Skylion007
Copy link
Collaborator

Also, curious if it's easy to add tests in our existing test framework for this

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@Skylion007 Skylion007 Skylion007 left review comments

@cpuhrsch cpuhrsch Awaiting requested review from cpuhrsch

Assignees
No one assigned
Labels
open source topic: not user facing topic category triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

Mask in MaskedTensor does not change device
4 participants
@zeshengzong @Skylion007 @pytorchbot @mikaylagawarecki
0