RaggedSoftMax

Applies the SoftMax function on an input tensor made of sequences across the sequence lengths specified, so that the values in the output lies in the range \([0,1]\), and the sum of all the values of each slice along the applied dimension is equal to 1. Values outside of the specified lengths are set to zero.

The SoftMax function is defined as:

\[\large{Softmax(x_{a_0,...,a_i,...,a_n}) = \frac{e^{x_{a_0,...,a_i,...,a_n}}}{\sum\limits_{j \in [:,...,a_i,...,:]}{e^{x_j}}}}\]

Inputs

input: tensor of type T1

bounds: tensor of type T2 describing the sequences lengths.

Outputs

output: tensor of type T1

Data Types

T1: float32

T2: int32

Shape Information

input and output are tensors with a shape of \([a_0,...,a_n]\) where \(n\in(2,3)\).

Volume Limits

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

Examples

RaggedSoftMax

in1 = network.add_input("input1", dtype=trt.float32, shape=(1, 2, 5))
bounds = network.add_input("bounds", dtype=trt.int32, shape=(1, 2, 1))
layer = network.add_ragged_softmax(in1, bounds)
network.mark_output(layer.get_output(0))

inputs[in1.name] = np.array(
    [
        [1.0, 1.0, 1.0, 1.0, 0.5],
        [3.0, 4.0, -20.0, 10.0, 5.0],
    ]
)
inputs[bounds.name] = np.array([5, 3])
outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array(
    [
        [
            [0.21, 0.21, 0.21, 0.21, 0.13],
            [0.26, 0.73, 0.01, 0.0, 0.0],
        ]
    ],
)

C++ API

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

Python API

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