mars.tensor.extract#

mars.tensor.extract(condition, a)[source]#

Return the elements of a tensor that satisfy some condition.

This is equivalent to mt.compress(ravel(condition), ravel(arr)). If condition is boolean mt.extract is equivalent to arr[condition].

Note that place does the exact opposite of extract.

Parameters

  • condition (array_like) – An array whose nonzero or True entries indicate the elements of arr to extract.

  • a (array_like) – Input tensor of the same size as condition.

Returns

extract – Rank 1 tensor of values from arr where condition is True.

Return type

Tensor

See also

take, put, copyto, compress, place

Examples

>>> import mars.tensor as mt

>>> arr = mt.arange(12).reshape((3, 4))
>>> arr.execute()
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> condition = mt.mod(arr, 3)==0
>>> condition.execute()
array([[ True, False, False,  True],
       [False, False,  True, False],
       [False,  True, False, False]])
>>> mt.extract(condition, arr).execute()
array([0, 3, 6, 9])

If condition is boolean:

>>> arr[condition].execute()
array([0, 3, 6, 9])