NonZero

计算输入张量中非零值的索引。返回的索引按行优先顺序排列。

输入

input: 类型为 T 的张量。

输出

output: 类型为 int32 的张量。

数据类型

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

形状信息

output 是形状为 \([D,C]\) 的张量，其中 \(D\) 是输入的维度数，\(C\) 是非零值的数量。

体积限制

input 和 output 最多可以有 \(2^{31}-1\) 个元素。

示例

非零 (NonZero)

opt_profile = get_runner.builder.create_optimization_profile()
get_runner.config.add_optimization_profile(opt_profile)

in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 3))
layer = network.add_non_zero(input=in1)
network.mark_output(layer.get_output(0))

inputs[in1.name] = np.array([[-3.0, -2.0, 0], [0.0, 1.0, 2.0]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape

expected[layer.get_output(0).name] = np.array(
    [
        [0, 0, 1, 1],
        [0, 1, 1, 2],
    ]
)

C++ API

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

Python API

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