Skip to content

Fix bool type indices #1023

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

Open
wants to merge 11 commits into
base: master
Choose a base branch
from
+170 −47
Open

Fix bool type indices #1023

Changes from 2 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
5051701
Fix bool type indices
botbigeyes Oct 30, 2025
8a72bed
Fix some errors
botbigeyes Nov 4, 2025
5e655c8
Add tests for bool type
botbigeyes Nov 5, 2025
5ce6e85
Add tests for bool type
botbigeyes Nov 5, 2025
bb9cd23
Add tests for bool type
botbigeyes Nov 6, 2025
a08ca6a
Add bool indices cases to benchmark
botbigeyes Nov 11, 2025
e5dbf7c
Merge branch 'master' into index_put
botbigeyes Nov 12, 2025
c1eacca
fix errors
botbigeyes Nov 12, 2025
effda03
fix errors
botbigeyes Nov 12, 2025
0609c5e
Merge remote-tracking branch 'origin/master' into index_put
botbigeyes Nov 17, 2025
ea8de52
The random seed is fixed
botbigeyes Nov 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions src/flag_gems/ops/index_put.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -255,6 +255,21 @@ def index_put(inp, indices, values, accumulate=False):
logger.debug("GEMS INDEX PUT")

indices = list(indices)
if len(indices) == 1 and indices[0].dtype == torch.bool:
mask = indices[0]

if mask.device != inp.device:
mask = mask.to(inp.device)

indices = list(torch.where(mask))

K = indices[0].numel()

if values.numel() == 1:
values = torch.full((K,), values.item(), dtype=inp.dtype, device=inp.device)
elif values.numel() == K:
values = values.reshape((K,))

indices = [
index.to(inp.device) if index.device != inp.device else index
for index in indices
Expand All @@ -275,6 +290,21 @@ def index_put_(inp, indices, values, accumulate=False):
logger.debug("GEMS INDEX PUT_")

indices = list(indices)
if len(indices) == 1 and indices[0].dtype == torch.bool:
mask = indices[0]

if mask.device != inp.device:
mask = mask.to(inp.device)

indices = list(torch.where(mask))

K = indices[0].numel()

if values.numel() == 1:
values = torch.full((K,), values.item(), dtype=inp.dtype, device=inp.device)
elif values.numel() == K:
values = values.reshape((K,))

indices = [
index.to(inp.device) if index.device != inp.device else index
for index in indices
Expand Down
Loading