mars.tensor.where#
- mars.tensor.where(condition, x=None, y=None)[source]#
Return elements, either from x or y, depending on condition.
If only condition is given, return
condition.nonzero()
.- Parameters
condition (array_like, bool) – When True, yield x, otherwise yield y.
x (array_like, optional) – Values from which to choose. x, y and condition need to be broadcastable to some shape.
y (array_like, optional) – Values from which to choose. x, y and condition need to be broadcastable to some shape.
- Returns
out – If both x and y are specified, the output tensor contains elements of x where condition is True, and elements from y elsewhere.
If only condition is given, return the tuple
condition.nonzero()
, the indices where condition is True.- Return type
Tensor or tuple of Tensors
Notes
If x and y are given and input arrays are 1-D, where is equivalent to:
[xv if c else yv for (c,xv,yv) in zip(condition,x,y)]
Examples
>>> import mars.tensor as mt
>>> mt.where([[True, False], [True, True]], ... [[1, 2], [3, 4]], ... [[9, 8], [7, 6]]).execute() array([[1, 8], [3, 4]])
>>> mt.where([[0, 1], [1, 0]]).execute() (array([0, 1]), array([1, 0]))
>>> x = mt.arange(9.).reshape(3, 3) >>> mt.where( x > 5 ).execute() (array([2, 2, 2]), array([0, 1, 2])) >>> mt.where(x < 5, x, -1).execute() # Note: broadcasting. array([[ 0., 1., 2.], [ 3., 4., -1.], [-1., -1., -1.]])
Find the indices of elements of x that are in goodvalues.
>>> goodvalues = [3, 4, 7] >>> ix = mt.isin(x, goodvalues) >>> ix.execute() array([[False, False, False], [ True, True, False], [False, True, False]]) >>> mt.where(ix).execute() (array([1, 1, 2]), array([0, 1, 1]))