diff --git a/heat/core/dndarray.py b/heat/core/dndarray.py index 1155ff188f..42d8d76f2e 100644 --- a/heat/core/dndarray.py +++ b/heat/core/dndarray.py @@ -921,7 +921,7 @@ def __getitem__(self, key: int | slice[int | None] | tuple[int, ...] | list[int] # TODO: remove this resplit!! key = manipulations.resplit(key) if key.larray.dtype in [torch.bool, torch.uint8]: - key = indexing.nonzero(key) + key = indexing.nonzero(key, as_tuple=False) if key.ndim > 1: key = list(key.larray.split(1, dim=1)) @@ -1631,7 +1631,7 @@ def __setitem__( to be used.""" key = manipulations.resplit(key) if key.larray.dtype in [torch.bool, torch.uint8]: - key = indexing.nonzero(key) + key = indexing.nonzero(key, as_tuple=False) if key.ndim > 1: key = list(key.larray.split(1, dim=1)) diff --git a/heat/core/indexing.py b/heat/core/indexing.py index 916aa450df..dd66646990 100644 --- a/heat/core/indexing.py +++ b/heat/core/indexing.py @@ -3,22 +3,23 @@ """ import torch -from typing import List, Dict, Any, TypeVar, Union, Tuple, Sequence from .communication import MPI from .dndarray import DNDarray -from . import sanitation +from . import factories +from .sanitation import sanitize_in from . import types +from . import manipulations __all__ = ["nonzero", "where"] -def nonzero(x: DNDarray) -> DNDarray: +def nonzero(x: DNDarray, as_tuple: bool = True) -> tuple[DNDarray, ...] | DNDarray: """ - Return a :class:`~heat.core.dndarray.DNDarray` containing the indices of the elements that are non-zero (using ``torch.nonzero``). - If ``x`` is split then the result is split in the first dimension. However, this :class:`~heat.core.dndarray.DNDarray` + Return a tuple of :class:`~heat.core.dndarray.DNDarray`s, one for each dimension of ``x``, + containing the indices of the non-zero elements in that dimension. If ``x`` is split then + the result is split in the first dimension. However, this :class:`~heat.core.dndarray.DNDarray` can be UNBALANCED as it contains the indices of the non-zero elements on each node. - Returns an array with one entry for each dimension of ``x``, containing the indices of the non-zero elements in that dimension. The values in ``x`` are always tested and returned in row-major, C-style order. The corresponding non-zero values can be obtained with: ``x[nonzero(x)]``. @@ -26,16 +27,16 @@ def nonzero(x: DNDarray) -> DNDarray: ---------- x: DNDarray Input array + as_tuple: bool, optional + Default is True for numpy-style nonzero output. If False, the output is a torch-style single 2D ``DNDarray`` of shape `(num_nonzero, ndim)` containing the indices of the non-zero elements. Examples -------- >>> import heat as ht >>> x = ht.array([[3, 0, 0], [0, 4, 1], [0, 6, 0]], split=0) >>> ht.nonzero(x) - DNDarray([[0, 0], - [1, 1], - [1, 2], - [2, 1]], dtype=ht.int64, device=cpu:0, split=0) + (DNDarray([0, 1, 1, 2], dtype=ht.int64, device=cpu:0, split=None), + DNDarray([0, 1, 2, 1], dtype=ht.int64, device=cpu:0, split=None)) >>> y = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=0) >>> y > 3 DNDarray([[False, False, False], @@ -48,48 +49,104 @@ def nonzero(x: DNDarray) -> DNDarray: [2, 0], [2, 1], [2, 2]], dtype=ht.int64, device=cpu:0, split=0) + (DNDarray([1, 1, 1, 2, 2, 2], dtype=ht.int64, device=cpu:0, split=None), + DNDarray([0, 1, 2, 0, 1, 2], dtype=ht.int64, device=cpu:0, split=None)) >>> y[ht.nonzero(y > 3)] DNDarray([4, 5, 6, 7, 8, 9], dtype=ht.int64, device=cpu:0, split=0) """ - sanitation.sanitize_in(x) - + sanitize_in(x) + + if not x.is_distributed(): + # nonzero indices as tuple + nonzero = torch.nonzero(input=x.larray, as_tuple=as_tuple) + # bookkeeping for final DNDarray construct + if as_tuple: + nonzero = list(nonzero) + for i, nz_tensor in enumerate(nonzero): + nonzero[i] = factories.array(nz_tensor, device=x.device, comm=x.comm) + return tuple(nonzero) + else: + # nonzero indices as single 2D DNDarray + return factories.array(nonzero, device=x.device, comm=x.comm) + + # distributed case lcl_nonzero = torch.nonzero(input=x.larray, as_tuple=False) - - # add offsets mapping from local indices to global indices if x is split - if x.split is not None: - _, _, slices = x.comm.chunk(x.shape, x.split) - lcl_nonzero[..., x.split] += slices[x.split].start - - if x.ndim == 1: - lcl_nonzero = lcl_nonzero.squeeze(dim=1) - - # compute global shape of the index array - gout = list(lcl_nonzero.shape) - if x.split is None: - is_split = None + nonzero_size = torch.tensor(lcl_nonzero.shape[0], dtype=torch.int64, device="cpu") + nonzero_dtype = types.canonical_heat_type(lcl_nonzero.dtype) + + # global nonzero_size + x.comm.Allreduce(MPI.IN_PLACE, nonzero_size, MPI.SUM) + # correct indices along split axis + _, displs = x.counts_displs() + lcl_nonzero[:, x.split] += displs[x.comm.rank] + + if x.split == 0: + # for split=0, the local nonzero indices are already globally ordered along the split axis + if as_tuple: # return indices as tuple of 1D DNDarrays + lcl_nonzero = lcl_nonzero.unbind(dim=1) + return tuple( + DNDarray( + nz_tensor, + gshape=(nonzero_size.item(),), + dtype=nonzero_dtype, + split=0, + device=x.device, + comm=x.comm, + balanced=False, + ) + for nz_tensor in lcl_nonzero + ) + else: # return indices as single 2D DNDarray + return DNDarray( + lcl_nonzero, + gshape=(nonzero_size.item(), x.ndim), + dtype=nonzero_dtype, + split=0, + device=x.device, + comm=x.comm, + balanced=False, + ) else: - gout[0] = x.comm.allreduce(gout[0], MPI.SUM) - is_split = 0 - - return DNDarray( - lcl_nonzero, - gshape=tuple(gout), - dtype=types.canonical_heat_type(lcl_nonzero.dtype), - split=is_split, - device=x.device, - comm=x.comm, - balanced=False, - ) - - -DNDarray.nonzero = lambda self: nonzero(self) + # construct global 2D DNDarray of nz indices: + shape_2d = (nonzero_size.item(), x.ndim) + global_nonzero = DNDarray( + lcl_nonzero, + gshape=shape_2d, + dtype=nonzero_dtype, + split=0, + device=x.device, + comm=x.comm, + balanced=False, + ) + # vectorized sorting of nz indices along axis 0 + global_nonzero.balance_() + global_nonzero = manipulations.vectorized_sort(global_nonzero, axis=0) + if as_tuple: # return indices as tuple of 1D DNDarrays + lcl_nonzero = global_nonzero.larray.unbind(dim=1) + return tuple( + DNDarray( + nz_tensor, + gshape=(nonzero_size.item(),), + dtype=nonzero_dtype, + split=0, + device=x.device, + comm=x.comm, + balanced=True, + ) + for nz_tensor in lcl_nonzero + ) + else: # return indices as single 2D DNDarray + return global_nonzero + + +DNDarray.nonzero = lambda self: nonzero(self, as_tuple=True) DNDarray.nonzero.__doc__ = nonzero.__doc__ def where( cond: DNDarray, - x: Union[None, int, float, DNDarray] = None, - y: Union[None, int, float, DNDarray] = None, + x: None | int | float | DNDarray = None, + y: None | int | float | DNDarray = None, ) -> DNDarray: """ Return a :class:`~heat.core.dndarray.DNDarray` containing elements chosen from ``x`` or ``y`` depending on condition. @@ -114,34 +171,42 @@ def where( Notes ----- - When only condition is provided, this function is a shorthand for :func:`nonzero`. + When only condition is provided, this function is a shorthand for :func:`nonzero` and the function returns a tuple + of :class:`~heat.core.dndarray.DNDarray`, analogously to ``numpy.where``. Examples -------- >>> import heat as ht >>> x = ht.arange(10, split=0) >>> ht.where(x < 5, x, 10 * x) - DNDarray([ 0, 1, 2, 3, 4, 50, 60, 70, 80, 90], dtype=ht.int64, device=cpu:0, split=0) + DNDarray(MPI-rank: 0, Shape: (10,), Split: 0, Local Shape: (10,), Device: cpu:0, Dtype: int32, Data: + [ 0, 1, 2, 3, 4, 50, 60, 70, 80, 90]) >>> y = ht.array([[0, 1, 2], [0, 2, 4], [0, 3, 6]]) >>> ht.where(y < 4, y, -1) - DNDarray([[ 0, 1, 2], - [ 0, 2, -1], - [ 0, 3, -1]], dtype=ht.int64, device=cpu:0, split=None) + DNDarray(MPI-rank: 0, Shape: (3, 3), Split: None, Local Shape: (3, 3), Device: cpu:0, Dtype: int64, Data: + [[ 0, 1, 2], + [ 0, 2, -1], + [ 0, 3, -1]]) """ - if cond.split is not None and (isinstance(x, DNDarray) or isinstance(y, DNDarray)): - if (isinstance(x, DNDarray) and cond.split != x.split) or ( - isinstance(y, DNDarray) and cond.split != y.split - ): - if len(y.shape) >= 1 and y.shape[0] > 1: - raise NotImplementedError("binary op not implemented for different split axes") + # binary where(cond, x, y) branch + if cond.split is not None and isinstance(y, DNDarray) and len(y.shape) >= 1 and y.shape[0] > 1: + if (isinstance(x, DNDarray) and cond.split != x.split) or cond.split != y.split: + raise NotImplementedError("binary op not implemented for different split axes") + if isinstance(x, (DNDarray, int, float)) and isinstance(y, (DNDarray, int, float)): + # Simple elementwise selection using arithmetic: + # cond == 0 -> take y, cond == 1 -> take x for var in [x, y]: if isinstance(var, int): var = float(var) return cond.dtype(cond == 0) * y + cond * x - elif x is None and y is None: - return nonzero(cond) + + # where(cond) "indices only" branch + elif x is None and y is None: # delegate to nonzero(cond) + return nonzero(cond) # tuple of DNDarrays, one per dimension + else: raise TypeError( - f"either both or neither x and y must be given and both must be DNDarrays or numerical scalars({type(x)}, {type(y)})" + "either both or neither x and y must be given and both must be " + f"DNDarrays or numerical scalars (got {type(x)}, {type(y)})" ) diff --git a/heat/core/linalg/eigh.py b/heat/core/linalg/eigh.py index 7f66212955..92cb904c02 100644 --- a/heat/core/linalg/eigh.py +++ b/heat/core/linalg/eigh.py @@ -74,7 +74,7 @@ def _subspaceiteration( device=columnnorms.device, ) * statistics.percentile(columnnorms, 100.0 * (1 - (k + safetyparam) / columnnorms.shape[0])) - ) + )[0] X = C[:, idx].balance() # actual subspace iteration diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py index b80faac41c..9702f3452e 100644 --- a/heat/core/manipulations.py +++ b/heat/core/manipulations.py @@ -53,6 +53,7 @@ "row_stack", "shape", "sort", + "vectorized_sort", "split", "squeeze", "stack", @@ -2899,6 +2900,204 @@ def sort( return tensor +def vectorized_sort( + a: DNDarray, + axis: int = -1, + stable: bool = True, + descending: bool = False, + resplit_result: bool = True, + return_sort_indices_only: bool = False, +) -> DNDarray: + """ + Performs a lexicographical sort along the specified axis. + + The array is transposed into an MxN matrix, where M is the + number of elements along the target `axis`, and N is the product of all + remaining dimensions. The lexicographical sorting prioritizes the leftmost + columns first, which acts as the primary sort key, with subsequent + columns acting as secondary, tertiary, etc., keys. + + Parameters + ---------- + a : DNDarray + The array to be sorted. + axis : int, optional + The axis along which to sort. If the split dimension of the array does + not match this axis, the array is resplit. + Default is -1 (last axis). + stable : bool, optional + Whether the sorting algorithm should be stable. Default is True. + descending : bool, optional + Whether to sort in descending order. Default is False. + resplit_result : bool, optional + Whether to resplit the final sorted array back to the original split + axis of the input array after rows are distributed. Default is True. + return_sort_indices_only : bool, optional + If True, bypasses the row exchange and returns only the global sort indices. Default is False. + + Returns + ------- + DNDarray + Either the final sorted array, or if `return_sort_indices_only` is True, the global sort indices. + """ + sanitation.sanitize_in(a) + + if not isinstance(axis, int): + raise ValueError("'axis' must be an int.") + if not isinstance(stable, bool): + raise ValueError("'stable' must be a bool.") + if not isinstance(descending, bool): + raise ValueError("'descending' must be a bool.") + if not isinstance(resplit_result, bool): + raise ValueError("'resplit_result' must be a bool.") + if not isinstance(return_sort_indices_only, bool): + raise ValueError("'return_sort_indices_only' must be a bool.") + + if len(a.gshape) == 0: + raise ValueError("dndarray must have atleast one dimension.") + + if axis >= len(a.gshape) or (axis < 0 and abs(axis) > len(a.gshape)): + raise ValueError(f"'axis'={axis} does not exist for '{len(a.gshape)}' dimenions.") + + def _permute_indices(data, idx): + sort_idx = torch.argsort(data[idx], stable=stable, descending=descending) + return idx[sort_idx] + + if len(a.gshape) == 1: + arr, idx = sort(a, axis=axis, descending=descending, return_sort_indices=True) + if return_sort_indices_only: + return idx + return arr.resplit_(a.split) if resplit_result else arr + + if not a.is_distributed(): + data = a.larray.transpose(axis, 0) + shape = data.shape + + data = data.reshape(shape[0], -1) + indices = torch.arange(0, data.shape[0]) + for i in range(data.shape[-1] - 1, -1, -1): + indices = _permute_indices(data[:, i], indices) + + data = data.reshape(shape)[indices] + return factories.array(data.transpose(axis, 0), split=None) + original_split = a.split + if axis != a.split: + a = resplit(a, axis) + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + size = comm.Get_size() + + data = a.larray.transpose(axis, 0) + original_shape = data.shape + inner_shape = original_shape[1:] + + local_count = data.shape[0] + total_rows = sum(comm.allgather(local_count)) + block_length = np.prod(inner_shape) if len(inner_shape) > 0 else 1 + + mpi_type: MPI.Datatype = a.comm.mpi_type_of(data.dtype) + + if rank == 0: + send_counts = np.array(comm.gather(local_count), dtype=int) + send_displ = np.insert(np.cumsum(send_counts)[:-1], 0, 0) + + buffer = torch.empty((total_rows,), dtype=data.dtype) + recv_args = [buffer, send_counts, send_displ, mpi_type] + else: + comm.gather(local_count) + + buffer = None + recv_args = None + + def gather_column(flat_idx: int): + idx = np.unravel_index(flat_idx, inner_shape) + slice_tuple = (slice(None),) + tuple(idx) + + local_col = data[slice_tuple].contiguous() + comm.Gatherv(local_col, recv_args) + return buffer + + if rank == 0: + indices = torch.arange(0, total_rows, dtype=torch.int64) + else: + indices = torch.empty(total_rows, dtype=torch.int64) + + for i in range(block_length - 1, -1, -1): + if rank == 0: + indices = _permute_indices(gather_column(i), indices) + else: + gather_column(i) + + comm.Bcast(indices) + + if return_sort_indices_only: + return factories.array(indices, split=None) + + offset, _, _ = a.comm.chunk((total_rows,), split=0, rank=rank) + + rank_slices = [a.comm.chunk((total_rows,), split=0, rank=i)[-1][0] for i in range(size)] + + local_slice = rank_slices[rank] + + assert all([s.step is None for s in rank_slices]) # Sanity check + + send_counts = np.zeros(size, dtype=np.int64) + send_indices = [] + + for recv_rank, s in enumerate(rank_slices): + recv_indices = indices[s] + + mask = (recv_indices >= offset) & (recv_indices < rank_slices[rank].stop) + + local_indices = recv_indices[mask] - offset + + send_counts[recv_rank] += mask.sum() + send_indices.append(local_indices) + + recv_counts = np.zeros(size, dtype=np.int64) + recv_indices = [list() for _ in range(size)] + + rank_indices_mapping = np.empty((local_slice.stop - local_slice.start,), dtype=np.int64) + + for i, idx in enumerate(indices[local_slice]): + for src_rank, src_slice in enumerate(rank_slices): + if not (src_slice.start <= idx < src_slice.stop): + continue + recv_counts[src_rank] += 1 + recv_indices[src_rank].append(idx.item()) + rank_indices_mapping[i] = src_rank + break + else: + raise RuntimeError(f"Index could not be resolved to a rank. Info: {i}, {idx}") + + send_counts *= block_length + recv_counts *= block_length + + send_displ = np.insert(np.cumsum(send_counts)[:-1], 0, 0) + recv_displ = np.insert(np.cumsum(recv_counts)[:-1], 0, 0) + + send_data = data[torch.cat(send_indices).tolist()].reshape(-1).contiguous() + recv_buf = torch.empty((recv_counts.sum().item(),), dtype=data.dtype) + + comm.Alltoallv( + [send_data, send_counts, send_displ, mpi_type], + [recv_buf, recv_counts, recv_displ, mpi_type], + ) + + sort_idx = np.argsort(rank_indices_mapping, stable=True) + inv_sort_idx = np.empty_like(sort_idx) + inv_sort_idx[sort_idx] = np.arange(sort_idx.size) + + recv_buf = recv_buf.view(-1, *inner_shape)[inv_sort_idx] + + sorted_array = factories.array(recv_buf.transpose(0, axis), is_split=a.split) + + if original_split != a.split and resplit_result: + return resplit(sorted_array, original_split) + return sorted_array + + def split(x: DNDarray, indices_or_sections: Iterable, axis: int = 0) -> List[DNDarray, ...]: """ Split a DNDarray into multiple sub-DNDarrays. diff --git a/tests/core/test_indexing.py b/tests/core/test_indexing.py index 61dda3fa4f..30a0311b3b 100644 --- a/tests/core/test_indexing.py +++ b/tests/core/test_indexing.py @@ -1,87 +1,154 @@ +import pytest + import heat as ht from heat.testing.basic_test import TestCase +import torch +import numpy as np -class TestIndexing(TestCase): - def test_nonzero(self): - # cases to test: - # not split - a = ht.array([[1, 2, 3], [4, 5, 2], [7, 8, 9]], split=None) - cond = a > 3 - nz = ht.nonzero(cond) - self.assertEqual(nz.gshape, (5, 2)) - self.assertEqual(nz.dtype, ht.int64) - self.assertEqual(nz.split, None) +def compare_ht_where_to_numpy_where(ht_res, np_res): + if isinstance(np_res, tuple): + assert isinstance(ht_res, tuple) - # split - a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=1) - cond = a > 3 - nz = cond.nonzero() - self.assertEqual(nz.gshape, (6, 2)) - self.assertEqual(nz.dtype, ht.int64) - self.assertEqual(nz.split, 0) - a[nz] = 10.0 - self.assertEqual(ht.all(a[nz] == 10), 1) + assert len(ht_res) == len(np_res) + for _ht_res, _np_res in zip(ht_res, np_res): + assert ht.equal(_ht_res, ht.array(_np_res)) + + assert ht_res[0].dtype == ht.int64 + + else: + assert np.allclose(ht_res.shape, np_res.shape) + assert ht.equal(ht_res, ht.array(np_res)) + assert ht_res[0].dtype == ht.types.canonical_heat_type(np_res.dtype) + +@pytest.mark.parametrize('split', [None, 0, 1]) +@pytest.mark.parametrize('cond_type', ['mean', 'max']) +def test_nonzero(split, cond_type): + a = ht.random.random((2*ht.comm.size, 3*ht.comm.size, 4*ht.comm.size)) + if cond_type == 'mean': + cond = a > a.mean() / 2 + elif cond_type == 'max': + cond = a == a.max() + else: + raise NotImplementedError + + nz_as_tuple = ht.nonzero(cond, as_tuple=True) + nz_as_tuple_ref = np.nonzero(cond.numpy()) + for i in range(len(nz_as_tuple)): + assert nz_as_tuple[i].dtype == ht.int64 + assert np.allclose(nz_as_tuple[i].numpy(), nz_as_tuple_ref[i]) + + nz_no_tuple = ht.nonzero(cond, as_tuple=False) + nz_no_tuple_ref = torch.nonzero(cond.resplit(None), as_tuple=False) + assert nz_no_tuple.dtype == ht.int64 + assert np.allclose(nz_no_tuple.numpy(), nz_no_tuple_ref.numpy()) + + if cond_type == 'max': + assert len(cond[cond]) == 1 + for me in nz_as_tuple: + assert me.shape == (1,) + assert nz_no_tuple.shape == (1, a.ndim) +@pytest.mark.parametrize('split', [None, 0, 1]) +def test_where_against_numpy(split): + # no x and y + a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=split) + cond = a > 3 + ht_res = ht.where(cond) + np_res = np.where(cond.numpy()) + compare_ht_where_to_numpy_where(ht_res, np_res) + if split is None or not cond.is_distributed(): + assert ht_res[0].split == None + else: + assert ht_res[0].split == 0 + + # x and y DNDarray + a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=split, dtype=ht.float32) + b = -a + cond = a > 3 + ht_res = ht.where(cond, a, b) + np_res = np.where(cond.numpy(), a.numpy(), b.numpy()) + compare_ht_where_to_numpy_where(ht_res, np_res) + assert ht_res.split == split + + # x and y float + a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=split, dtype=ht.float32) + cond = a > 3 + ht_res = ht.where(cond, 1., -1.) + np_res = np.where(cond.numpy(), 1., -1.) + compare_ht_where_to_numpy_where(ht_res, np_res.astype(np.float32)) + assert ht_res.split == split + + # x and y int + a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=split, dtype=ht.float32) + cond = a > 3 + ht_res = ht.where(cond, 1, -1) + np_res = np.where(cond.numpy(), 1, -1) + compare_ht_where_to_numpy_where(ht_res, np_res.astype(np.int32)) + assert ht_res.split == split + + + +class TestIndexing(TestCase): + def test_nonzero_special_cases(self): # edge case: single non-zero element for split in [None, 0, 1]: a = ht.zeros((4, 3), dtype=ht.bool, split=split) a[1, 2] = True - nz = ht.indexing.nonzero(a) - a.resplit_(None) - nz.resplit_(None) - self.assertEqual(nz.gshape, (1, 2)) + nz = ht.indexing.nonzero(a, as_tuple=False) self.assertTrue(ht.allclose(a[nz], a[a])) + a.comm.Barrier() + # as_tuple = False (torch-style output) + a = ht.array([[1, 0, 0], [0, 4, 1], [0, 6, 0]], split=1) + nz = ht.nonzero(a, as_tuple=False) + self.assertEqual(nz.gshape, (4, 2)) + self.assertEqual(nz.dtype, ht.int64) + if a.is_distributed(): + self.assertEqual(nz.split, 0) + else: + self.assertEqual(nz.split, None) + t_a = a.resplit_(None).larray + t_nz = torch.nonzero(t_a, as_tuple=False) + self.assertTrue(ht.equal(nz, ht.array(t_nz))) - def test_where(self): - # cases to test - # no x and y - a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=None) - cond = a > 3 - wh = ht.where(cond) - self.assertEqual(wh.gshape, (6, 2)) - self.assertEqual(wh.dtype, ht.int64) - self.assertEqual(wh.split, None) - # split - a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=1) - cond = a > 3 - wh = ht.where(cond) - self.assertEqual(wh.gshape, (6, 2)) - self.assertEqual(wh.dtype, ht.int64) - self.assertEqual(wh.split, 0) + # attribute error + a = a.numpy() + with self.assertRaises(TypeError): + ht.nonzero(a) + def test_where_special_cases(self): # not split cond a = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, 4.0], [0.0, 3.0, 6.0]], split=None) res = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, -1.0], [0.0, 3.0, -1.0]], split=None) - wh = ht.where(a < 4.0, a, -1) + ht_res = ht.where(a < 4.0, a, -1) self.assertTrue( ht.equal(a[ht.nonzero(a < 4)], ht.array([0.0, 1.0, 2.0, 0.0, 2.0, 0.0, 3.0])) ) - self.assertTrue(ht.equal(wh, res)) - self.assertEqual(wh.gshape, (3, 3)) - self.assertEqual(wh.dtype, ht.float32) + self.assertTrue(ht.equal(ht_res, res)) + self.assertEqual(ht_res.gshape, (3, 3)) + self.assertEqual(ht_res.dtype, ht.float32) # split cond a = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, 4.0], [0.0, 3.0, 6.0]], split=0) res = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, -1.0], [0.0, 3.0, -1.0]], split=0) - wh = ht.where(a < 4.0, a, -1) - self.assertTrue(ht.all(wh[ht.nonzero(a >= 4)] == -1)) - self.assertTrue(ht.equal(wh, res)) - self.assertEqual(wh.gshape, (3, 3)) - self.assertEqual(wh.dtype, ht.float32) - self.assertEqual(wh.split, 0) + ht_res = ht.where(a < 4.0, a, -1) + self.assertTrue(ht.all(ht_res[ht.nonzero(a >= 4)] == -1)) + self.assertTrue(ht.equal(ht_res, res)) + self.assertEqual(ht_res.gshape, (3, 3)) + self.assertEqual(ht_res.dtype, ht.float32) + self.assertEqual(ht_res.split, 0) a = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, 4.0], [0.0, 3.0, 6.0]], split=1) res = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, -1.0], [0.0, 3.0, -1.0]], split=1) - wh = ht.where(a < 4.0, a, -1.0) - self.assertTrue(ht.equal(wh, res)) - self.assertEqual(wh.gshape, (3, 3)) - self.assertEqual(wh.dtype, ht.float) - self.assertEqual(wh.split, 1) + ht_res = ht.where(a < 4.0, a, -1.0) + self.assertTrue(ht.equal(ht_res, res)) + self.assertEqual(ht_res.gshape, (3, 3)) + self.assertEqual(ht_res.dtype, ht.float) + self.assertEqual(ht_res.split, 1) with self.assertRaises(TypeError): - ht.where(cond, a) + ht.where(a < 3, a) with self.assertRaises(NotImplementedError): - ht.where(cond, ht.ones((3, 3), split=0), ht.ones((3, 3), split=1)) + ht.where(a < 3, ht.ones((3, 3), split=0), ht.ones((3, 3), split=1))