NonZero

Computes the indices of the input tensor where the value is non-zero. The returned indices are in row-major order.

Inputs

input: tensor of type T.

Outputs

output: tensor of type int32.

Data Types

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

Shape Information

output is a tensor with a shape of \([D,C]\), where \(D\) is the number of dimensions of the input and \(C\) is the number of non-zero values.

Volume Limits

input and output can have up to \(2^{31}-1\) elements.

Examples

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

For more information about the C++ INonZeroLayer operator, refer to the C++ INonZeroLayer documentation.

Python API

For more information about the Python INonZeroLayer operator, refer to the Python INonZeroLayer documentation.