Select¶
Selects elements from true and false tensors based on a condition tensor into an output tensor. When applicable, broadcasting is used (refer to Shape Information for more information).
Outputs¶
output: tensor of type T.
Data Types¶
T: int32, int64, float16, float32, bfloat16, bool
Shape Information¶
condition, true and false tensors must have the same rank. For each dimension, their lengths must match, or one of them must be equal to 1. In the latter case, the tensor is broadcast along that axis.
The output tensor has the same rank as the inputs. For each output dimension, its length is equal to the lengths of the corresponding input dimensions if they match, otherwise it is equal to the length that is not 1.
Examples¶
Select
condition_input = network.add_input("condition", dtype=trt.bool, shape=(2, 5))
true_input = network.add_input("true", dtype=trt.float32, shape=(2, 5))
false_input = network.add_input("false", dtype=trt.float32, shape=(1, 5))
layer = network.add_select(condition_input, true_input, false_input)
network.mark_output(layer.get_output(0))
inputs[condition_input.name] = np.array([[True, True, False, False, True], [False, True, False, False, True]])
inputs[true_input.name] = np.array([[1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0]])
inputs[false_input.name] = np.array([[0.0, 0.0, 0.0, 0.0, 0.0]])
outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[1.0, 1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 0.0, 1.0]])
C++ API¶
For more information about the C++ ISelectLayer operator, refer to the C++ ISelectLayer documentation.
Python API¶
For more information about the Python ISelectLayer operator, refer to the Python ISelectLayer documentation.