Shuffle

Reshapes and transposes an input tensor into an output tensor.

The layer first transposes the tensor, then reshapes it, and finally, it applies another transpose on the tensor.

Attributes

first_transpose The permutation applied by the first transpose operation. Default: Identity Permutation.

reshape_dims The reshaped dimensions. The product of the dimensions must be equal to the product of the input dimensions. Two special values can be used as dimension value:

• 0 Copies the corresponding dimension from the input. If the number of reshape dimensions is less than the input, 0s are resolved by aligning the most significant input dimensions. See also: zero_is_placeholder.

• -1 Infers that particular dimension by looking at the input and the rest of the reshape dimensions. Only one dimension is permitted to be specified as -1. Avoid using -1 if the input can have zero volume and any of the other reshape dimensions can be zero (after resolving special treatment of 0), because the solution for the -1 becomes indeterminate and TensorRT will report an error.

second_transpose The permutation applied by the second transpose operation. Default: Identity Permutation.

zero_is_placeholder The meaning of 0 in the reshape dimensions. If true, a 0 in the reshape dimensions denotes copying the corresponding dimension from the first input tensor. If false, a 0 in the reshape dimensions represents a zero-length dimension. Default: true.

Inputs

input0: tensor of type T.

input1: optional tensor of type Int32 or Int64 with reshape_dims.

Outputs

output: tensor of type T.

Data Types

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

Shape Information

output is a tensor with rank n.

input1 has shape \([n]\).

Examples

Shuffle

in1 = network.add_input("input1", dtype=trt.float32, shape=(3, 4))
layer = network.add_shuffle(in1)
layer.first_transpose = trt.Permutation([1, 0])
layer.reshape_dims = trt.Dims([2, 6])
network.mark_output(layer.get_output(0))

inputs[in1.name] = 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],
    ]
)

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

expected[layer.get_output(0).name] = np.array(
    [[1.0, 10.0, 100.0, 2.0, 20.0, 200.0], [3.0, 30.0, 300.0, 4.0, 40.0, 400.0]]
)

Shuffle Reshape Infer

in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 3, 4))
layer = network.add_shuffle(in1)
layer.first_transpose = trt.Permutation([1, 0, 2])
layer.reshape_dims = trt.Dims([2, -1, 3])
network.mark_output(layer.get_output(0))

inputs[in1.name] = 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],
        ],
        [
            [5.0, 6.0, 7.0, 8.0],
            [50.0, 60.0, 70.0, 80.0],
            [500.0, 600.0, 700.0, 800.0],
        ],
    ]
)

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

expected[layer.get_output(0).name] = np.array(
    [
        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 10.0], [20.0, 30.0, 40.0]],
        [[50.0, 60.0, 70.0], [80.0, 100.0, 200.0], [300.0, 400.0, 500.0], [600.0, 700.0, 800.0]],
    ]
)

C++ API

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

Python API

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