mars.tensor.compress¶

mars.tensor.
compress
(condition, a, axis=None, out=None)[source]¶ Return selected slices of a tensor along given axis.
When working along a given axis, a slice along that axis is returned in output for each index where condition evaluates to True. When working on a 1D array, compress is equivalent to extract.
 Parameters
condition (1D tensor of bools) – Tensor that selects which entries to return. If len(condition) is less than the size of a along the given axis, then output is truncated to the length of the condition tensor.
a (array_like) – Tensor from which to extract a part.
axis (int, optional) – Axis along which to take slices. If None (default), work on the flattened tensor.
out (Tensor, optional) – Output tensor. Its type is preserved and it must be of the right shape to hold the output.
 Returns
compressed_array – A copy of a without the slices along axis for which condition is false.
 Return type
Tensor
See also
take
,choose
,diag
,diagonal
,select
Tensor.compress
Equivalent method in ndarray
mt.extract
Equivalent method when working on 1D arrays
Examples
>>> import mars.tensor as mt
>>> a = mt.array([[1, 2], [3, 4], [5, 6]]) >>> a.execute() array([[1, 2], [3, 4], [5, 6]]) >>> mt.compress([0, 1], a, axis=0).execute() array([[3, 4]]) >>> mt.compress([False, True, True], a, axis=0).execute() array([[3, 4], [5, 6]]) >>> mt.compress([False, True], a, axis=1).execute() array([[2], [4], [6]])
Working on the flattened tensor does not return slices along an axis but selects elements.
>>> mt.compress([False, True], a).execute() array([2])