SortOperation

Applicable Products

Hardware Model	Supported	Remarks
Atlas 350 accelerator cards	√	-
Atlas A3 inference products/Atlas A3 training products	√	-
Atlas A2 training products/Atlas A2 inference products	√	-
Atlas training products	√	The float data type is not supported.
Atlas inference products	√	The float and bf16 data types are not supported.
Atlas 200I/500 A2 inference products	√	The float data type is not supported.

Description

Implements postprocessing computation by sorting input tensors in descending order in the last dimension, retaining the maximum num elements, and outputting the sorted tensor and the index corresponding to each element.

Operator Context

Figure 1 Context of the SortOperation operator

Definition

struct SortParam {
    SVector<int32_t> num;
    uint8_t rsv[8] = {0};
};

Parameters

Member	Type	Default Value	Description
num	SVector<int32_t>	-	Maximum number of elements that is retained after sorting. num is an SVector containing only one value. The value must be greater than 0 and less than or equal to the size of the last dimension of input x.
rsv[8]	uint8_t	{0}	Reserved

Member

Type

Default Value

Description

num

SVector<int32_t>

Maximum number of elements that is retained after sorting.

num is an SVector containing only one value. The value must be greater than 0 and less than or equal to the size of the last dimension of input x.

rsv[8]

uint8_t

{0}

Reserved

Input

Parameter	Dimension	Data Type	Format	Description
x	[dim_0, dim_1, ..., dim_n]	float16/bf16/float	ND	The last dimension must contain at least num elements.

Output

Parameter	Dimension	Data Type	Format	Description
output	[dim_0, dim_1, ..., num]	float16/bf16/float	ND	Maximum num elements after sorting in the last dimension.
indices	[dim_0, dim_1, ..., num]	int32	ND	Original indexes corresponding to the maximum num elements.

Constraints

num is an SVector containing only one value. The value must be greater than 0 and less than or equal to the size of the last dimension of input x.
On Atlas inference products, num must be greater than or equal to 16, and the dimension of the non-tail axis must be less than or equal to 580000000.

API Calling Example

Input:

num = [1]
x = [[3.0, 4.0],
     [5.0, 6.0]]

Output:

output = [[4.0],
          [6.0]]
indices = [[1.0],
           [1.0]]

Parent topic: atb/infer_op_params.h