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.

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

input: tensor of type T.

Outputs

output: tensor of type T.

Data Types

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

Shape Information

input and output are tensors with rank 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.