Add MaskedTensor support to _is_any_true #128574
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/128574
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ✅ No FailuresAs of commit 4f0a867 with merge base b4a7b54 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
nowtryz
force-pushed
the
patch-masked-is-any-true
branch
from
d021ffb
to
264ef1c
Compare
nowtryz
force-pushed
the
patch-masked-is-any-true
branch
from
264ef1c
to
b9a203b
Compare
| data = _get_data(args[0]) | ||
| mask = _maybe_get_mask(args[0]) | ||
| if mask is None: | ||
| raise ValueError(f"__torch_dispatch__, {func}: expected a mask tensor") |
There was a problem hiding this comment.
| raise ValueError(f"__torch_dispatch__, {func}: expected a mask tensor") | |
| raise ValueError(f"__torch_dispatch__, {func}: expected args[0] to be a MaskedTensor") |
|
|
||
| if func(mask): | ||
| return MaskedTensor(func(data & mask), torch.tensor(True)) | ||
| return MaskedTensor(torch.tensor(0), torch.tensor(False)) |
There was a problem hiding this comment.
Could this function just be return func(data & mask)? e.g. does the return necessarily have to be a MaskedTensor?
If everything is masked out, technically "any" of an empty set is just False?
There was a problem hiding this comment.
The overload is supposed to return a masked tensor but it is true that if there is no specified true value, then the result is the just False. I'll still wrap the result in a MaskedTensor for coherence, otherwise MaskedTensor.__torch_function__ will still try to wrap the result in a MaskedTensor and fail.
There was a problem hiding this comment.
otherwise MaskedTensor.torch_function will still try to wrap the result in a MaskedTensor and fail.
Hmm I think its still worth thinking about, as we could simply change the __torch_function__ to relax the condition.
In theory a masked tensor with all True mask should behave as-if it were a plain tensor, but performance-wise maybe we prefer plain tensors?
There was a problem hiding this comment.
Nevermind, torch_function does not enforce subclass, I'll fix this today or tomorrow.
soulitzer
added
the
triaged
nowtryz
force-pushed
the
patch-masked-is-any-true
branch
from
b9a203b
to
4f0a867
Compare
Fixes #128557
If there is a better way to detect autograd anomalies consistently, feel free to share your ideas. This is a dirty check.