Squeeze

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

When axes is absent, removes every dimension of the input whose size is statically 1 in the network definition. Each dynamic dimension is retained and must not be 1 at runtime, because removing it would change the output rank, which is fixed at definition time. Zero-sized dimensions are retained: for example, an input of shape [0, 1, 3] produces [0, 3] when axes are absent.

Inputs

input0: tensor of type T.

input1: tensor of type Int32 or Int64. Optional.

Outputs

output: tensor of type T.

Data Types

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

Shape Information

When input1 is provided, it has shape \([n]\) and output is a tensor with rank of \(rank(input) - n\).

When input1 is absent, output has the rank of the input minus the number of input dimensions that are statically 1.

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

# Without an axes input, every dimension that is statically 1 in the network definition is removed.
in2 = network.add_input("input2", dtype=trt.float32, shape=(1, 3, 1, 4))
layer2 = network.add_squeeze(in2)
network.mark_output(layer2.get_output(0))

inputs[in2.name] = test_data.reshape(1, 3, 1, 4)
outputs[layer2.get_output(0).name] = layer2.get_output(0).shape
expected[layer2.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.