fix: support masked_scatter by lowering path and corner case of maske…

[chohk88]

Description

Implemented support for masked_scatter in the lowering path, referring to this implementation in PyTorch Inductor.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

• New feature (non-breaking change which adds functionality)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[apbose]

@chohk88 the implementation looks perfect as such.
Just a slight detail I was wondering. Torch mentions that it supports broadcasting between mask and self tensor, but I see this example working too, basically broadcasting between source and mask. is this supported?

input = torch.zeros(3, 3)
mask = torch.tensor([[1, 0, 1], [0, 0, 1], [0, 0, 0]], dtype=torch.bool)
source = torch.tensor([2, 3, 4])
out = input.masked_scatter_(mask, source)
print("out is", out)

[chohk88]

Thanks @apbose, great question! I checked PyTorch’s behavior and it:

1. Broadcasts only input and mask (not source)
2. Flattens source as-is
3. Uses only the first mask.sum() elements of source, ignoring any extras

For example, with source=[2,3,4,5] and three True mask positions, PyTorch applies [2,3,4] and drops the 5 without error.

Our current decomposition does exactly that (broadcasts input/mask, flattens source, then cumsum→gather→where), so no code changes are needed. Let me know if you spot any other edge cases!

[apbose]

great! Thanks for the detailed explanation

[pytorch-bot]

No ciflow labels are configured for this repo.
For information on how to enable CIFlow bot see this wiki