From 8e9342e9973238a94a8e9c962c0037a9d45d6b5b Mon Sep 17 00:00:00 2001 From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:32:27 +0200 Subject: [PATCH 1/4] Add failing test --- tests/core/test_sorting.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py index 49ccc11266..0d8aeeb91b 100644 --- a/tests/core/test_sorting.py +++ b/tests/core/test_sorting.py @@ -73,6 +73,7 @@ def test_vectorized_sort_multi_dim(self, orig_shape, split, axis, stable, descen assert np.isclose(res, expected_res).all() assert np.equal(sort_idx, res_idxs).all() + assert a.device == res.device @pytest.mark.parametrize("descending", [False, True]) @pytest.mark.parametrize("stable", [False, True]) From b5c86330d01db7df4a1633b3fa23d710c6c89311 Mon Sep 17 00:00:00 2001 From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:43:36 +0200 Subject: [PATCH 2/4] Fix test --- tests/core/test_sorting.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py index 0d8aeeb91b..ee68ffa412 100644 --- a/tests/core/test_sorting.py +++ b/tests/core/test_sorting.py @@ -68,12 +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_idx.device @pytest.mark.parametrize("descending", [False, True]) @pytest.mark.parametrize("stable", [False, True]) From cadff06d20a358b8b686bd4e42586f6f6d24d486 Mon Sep 17 00:00:00 2001 From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:50:13 +0200 Subject: [PATCH 3/4] Specify device in a few places --- heat/core/manipulations.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py index c24b3f6543..80404896e9 100644 --- a/heat/core/manipulations.py +++ b/heat/core/manipulations.py @@ -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 @@ -3004,7 +3004,7 @@ 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) @@ -3012,7 +3012,7 @@ def _permute_indices(data, idx): 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 @@ -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) @@ -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) @@ -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)) @@ -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) From 26580db2dd9db4d03056b2e2dff28eda6665824c Mon Sep 17 00:00:00 2001 From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:52:23 +0200 Subject: [PATCH 4/4] Fix test for real this time --- tests/core/test_sorting.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py index ee68ffa412..1a38ef1453 100644 --- a/tests/core/test_sorting.py +++ b/tests/core/test_sorting.py @@ -74,7 +74,7 @@ def test_vectorized_sort_multi_dim(self, orig_shape, split, axis, stable, descen 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_idx.device + assert a.device == res_idxs.device @pytest.mark.parametrize("descending", [False, True]) @pytest.mark.parametrize("stable", [False, True])