Skip to content

[Bug fix] Return vectorized_sort result on the same device as input #2512

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

Merged
Berkant03 merged 5 commits into from
Sep 1, 2026
Merged

[Bug fix] Return vectorized_sort result on the same device as input #2512

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
8e9342e
Add failing test
brownbaerchen Aug 31, 2026
b5c8633
Fix test
brownbaerchen Aug 31, 2026
cadff06
Specify device in a few places
brownbaerchen Aug 31, 2026
26580db
Fix test for real this time
brownbaerchen Aug 31, 2026
f7ac78d
Merge branch 'main' into vectorized_sort_gpu
brownbaerchen Sep 1, 2026
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
18 changes: 10 additions & 8 deletions heat/core/manipulations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2977,10 +2977,10 @@ def _permute_indices(data, idx):
indices = _permute_indices(local_data[:, i], indices)

if return_sort_indices_instead:
return factories.array(indices, split=None)
return factories.array(indices, split=None, device=a.device)

local_data = local_data.reshape(shape)[indices].transpose(axis, 0)
return factories.array(local_data, split=None)
return factories.array(local_data, split=None, device=a.device)

# distributed vectorized sort
original_split = a.split
Expand All @@ -3004,15 +3004,15 @@ def _permute_indices(data, idx):
total_rows = a.gshape[axis]
block_length = np.prod(inner_shape)

send_buf = torch.tensor([local_count], dtype=torch.int64)
send_buf = torch.tensor([local_count], dtype=torch.int64, device=local_data.device)
local_counts = torch.empty(size, dtype=torch.int64)
comm.Gather(send_buf, local_counts, root=0)

if rank == 0:
send_counts = local_counts.numpy()
send_displ = np.insert(np.cumsum(send_counts)[:-1], 0, 0)

buffer = torch.empty((total_rows,), dtype=local_data.dtype)
buffer = torch.empty((total_rows,), dtype=local_data.dtype, device=local_data.device)
recv_args = (buffer, send_counts, send_displ)
else:
buffer = None
Expand All @@ -3026,7 +3026,7 @@ def _gather_column(flat_idx: int):
comm.Gatherv(local_col, recv_args, root=0)
return buffer

indices = torch.arange(0, total_rows, dtype=torch.int64)
indices = torch.arange(0, total_rows, dtype=torch.int64, device=local_data.device)

for i in range(block_length - 1, -1, -1):
buffer = _gather_column(i)
Expand All @@ -3036,7 +3036,7 @@ def _gather_column(flat_idx: int):
comm.Bcast(indices, root=0)

if return_sort_indices_instead:
return factories.array(indices, split=None)
return factories.array(indices, split=None, device=a.device)

offset, _, _ = comm.chunk((total_rows,), split=0, rank=rank)

Expand Down Expand Up @@ -3082,7 +3082,9 @@ def _gather_column(flat_idx: int):
recv_displ = np.insert(np.cumsum(recv_counts)[:-1], 0, 0)

send_data = local_data[torch.cat(send_indices).tolist()].reshape(-1).contiguous()
recv_buf = torch.empty((recv_counts.sum().item(),), dtype=local_data.dtype)
recv_buf = torch.empty(
(recv_counts.sum().item(),), dtype=local_data.dtype, device=local_data.device
)

comm.Alltoallv((send_data, send_counts, send_displ), (recv_buf, recv_counts, recv_displ))

Expand All @@ -3095,7 +3097,7 @@ def _gather_column(flat_idx: int):
if is_1d:
recv_buf = recv_buf.squeeze(-1)

sorted_array = factories.array(recv_buf.transpose(0, axis), is_split=a.split)
sorted_array = factories.array(recv_buf.transpose(0, axis), is_split=a.split, device=a.device)

if original_split != a.split and resplit_result:
return resplit(sorted_array, original_split)
Expand Down
10 changes: 6 additions & 4 deletions tests/core/test_sorting.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,11 +68,13 @@ def test_vectorized_sort_multi_dim(self, orig_shape, split, axis, stable, descen
sort_idx = np.lexsort(keys)
expected_res = arr[sort_idx].reshape(shape).swapaxes(0, axis)

res = ht.vectorized_sort(a, axis=axis, stable=stable, descending=descending).numpy()
res_idxs = ht.vectorized_sort(a, axis=axis, stable=stable, descending=descending, return_sort_indices_instead=True).numpy()
res = ht.vectorized_sort(a, axis=axis, stable=stable, descending=descending)
res_idxs = ht.vectorized_sort(a, axis=axis, stable=stable, descending=descending, return_sort_indices_instead=True)

assert np.isclose(res, expected_res).all()
assert np.equal(sort_idx, res_idxs).all()
assert np.isclose(res.numpy(), expected_res).all()
assert a.device == res.device
assert np.equal(sort_idx, res_idxs.numpy()).all()
assert a.device == res_idxs.device

@pytest.mark.parametrize("descending", [False, True])
@pytest.mark.parametrize("stable", [False, True])
Expand Down
Loading