Squeeze

Reshapes the input tensor by removing dimensions specified by axes. Corresponding dimensions must have length 1.

Inputs

input0: tensor of type T.

input1: tensor of type Int32 or Int64

Outputs

output: tensor of type T.

Data Types

T: bool, int4, int8, int32, int64, float8, float16, float32, bfloat16

Shape Information

input1 has shape \([n]\).

output is a tensor with rank of \(rank(input) - n\).

Examples

Squeeze

in1 = network.add_input("input1", dtype=trt.float32, shape=(3, 1, 4, 1))
axes_weights = trt.Weights(np.array([1, -1], dtype=np.int64))
axes_layer = network.add_constant((2,), axes_weights)
axes_tensor = axes_layer.get_output(0)
layer = network.add_squeeze(in1, axes_tensor)
network.mark_output(layer.get_output(0))

test_data = np.array(
    [
        [1.0, 2.0, 3.0, 4.0],
        [10.0, 20.0, 30.0, 40.0],
        [100.0, 200.0, 300.0, 400.0],
    ]
)

inputs[in1.name] = test_data.reshape(3, 1, 4, 1)

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

expected[layer.get_output(0).name] = test_data

C++ API

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

Python API

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