选择

基于**condition**张量，从 **thenInput** 和 **elseInput** 张量中选择元素，放入输出张量。 如果适用，则使用广播（有关更多信息，请参阅形状信息）。

输入

condition: `bool` 类型的张量。

thenInput: `T` 类型的张量。

elseInput: `T` 类型的张量。

输出

output: `T` 类型的张量。

数据类型

T: int32, int64, float16, float32, bfloat16, bool

形状信息

condition、thenInput 和 elseInput 张量必须具有相同的秩。 对于每个维度，它们的长度必须匹配，或者其中一个必须等于 1。 在后一种情况下，张量将沿该轴广播。

output 张量与输入张量具有相同的秩。 对于每个输出维度，如果它们匹配，则其长度等于相应输入维度的长度，否则等于不为 1 的长度。

体积限制

每个输入和输出最多可以有 \(2^{31}-1\) 个元素。

示例

选择

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

有关 C++ ISelectLayer 算子的更多信息，请参阅 C++ ISelectLayer 文档。

Python API

有关 Python ISelectLayer 算子的更多信息，请参阅 Python ISelectLayer 文档。