where¶

paddle.fluid.layers.nn. where ( condition ) [source]

Return an int64 tensor with rank 2, specifying the coordinate of true element in condition.

Parameters

condition (Variable) – A bool tensor with rank at least 1, the data type is bool.

Returns

The tensor variable storing a 2-D tensor, which involves all coordinate.

Return type

Variable, the output data type is int64.

Examples

import paddle.fluid as fluid
import paddle.fluid.layers as layers
import numpy as np

# condition is a tensor [True, False, True]
condition = layers.assign(np.array([1, 0, 1], dtype='int32'))
condition = layers.cast(condition, 'bool')
out = layers.where(condition) # [[0], [2]]

# condition is a tensor [[True, False], [False, True]]
condition = layers.assign(np.array([[1, 0], [0, 1]], dtype='int32'))
condition = layers.cast(condition, 'bool')
out = layers.where(condition) # [[0, 0], [1, 1]]

# condition is a tensor [False, False, False]
condition = layers.assign(np.array([0, 0, 0], dtype='int32'))
condition = layers.cast(condition, 'bool')
out = layers.where(condition) # [[]]