Edit

DML_TOP_K_OPERATOR_DESC structure (directml.h)

  • Article

Selects the largest K elements from each sequence along an axis of the InputTensor, and returns the values and indices of those elements in the OutputValueTensor and OutputIndexTensor, respectively. A sequence refers to one of the sets of elements that exist along the Axis dimension of the InputTensor.

Syntax

struct DML_TOP_K_OPERATOR_DESC {
  const DML_TENSOR_DESC *InputTensor;
  const DML_TENSOR_DESC *OutputValueTensor;
  const DML_TENSOR_DESC *OutputIndexTensor;
  UINT                  Axis;
  UINT                  K;
};

Members

InputTensor

Type: const DML_TENSOR_DESC*

The input tensor containing elements to select.

OutputValueTensor

Type: const DML_TENSOR_DESC*

The output tensor to write the values of the top K elements to. This tensor must have sizes equal to the InputTensor, except for the dimension specified by the Axis parameter, which must have a size equal to K.

The K values selected from each input sequence are guaranteed to be sorted descending (largest to smallest).

OutputIndexTensor

Type: const DML_TENSOR_DESC*

The output tensor to write the indices of the top K elements to. This tensor must have sizes equal to the InputTensor, except for the dimension specified by the Axis parameter, which must have a size equal to K.

The indices returned in this tensor are measured relative to the beginning of their sequence (as opposed to the beginning of the tensor). For example, an index of 0 always refers to the first element for all sequences in an axis.

In cases where two or more elements in the top-K have the same value (that is, when there is a tie), the indices of both elements are included, and are guaranteed to be ordered by ascending element index.

Axis

Type: UINT

The index of the dimension to select elements across. This value must be less than the DimensionCount of the InputTensor.

K

Type: UINT

The number of elements to select. K must be greater than 0, but less than the number of elements in the InputTensor along the dimension specified by Axis.

Examples

Example 1

InputTensor: (Sizes:{1,1,3,4}, DataType:FLOAT32)
[[[[ 0,  1, 10, 11],
   [ 3,  2,  9,  8],
   [ 4,  5,  6,  7]]]]

Axis: 3
K:    2
   
OutputValueTensor: (Sizes:{1,1,3,2}, DataType:FLOAT32)
[[[[11, 10],
   [ 9,  8],
   [ 7,  6]]]]

OutputIndexTensor: (Sizes:{1,1,3,2}, DataType:UINT32)
[[[[3, 2],
   [2, 3],
   [3, 2]]]]

Example 2. Using a different axis

InputTensor: (Sizes:{1,1,3,4}, DataType:FLOAT32)
[[[[ 0,  1, 10, 11],
   [ 3,  2,  9,  8],
   [ 4,  5,  6,  7]]]]

Axis: 2
K:    2
   
OutputValueTensor: (Sizes:{1,1,2,4}, DataType:FLOAT32)
[[[[ 4,  5, 10, 11],
   [ 3,  2,  9,  8]]]]

OutputIndexTensor: (Sizes:{1,1,2,4}, DataType:UINT32)
[[[[2, 2, 0, 0],
   [1, 1, 1, 1]]]]

Example 3. Tied values

InputTensor: (Sizes:{1,1,3,4}, DataType:FLOAT32)
[[[[1, 2, 2, 3],
   [3, 4, 5, 5],
   [6, 6, 6, 6]]]]

Axis: 3
K:    3
   
OutputValueTensor: (Sizes:{1,1,3,3}, DataType:FLOAT32)
[[[[3, 2, 2],
   [5, 5, 4],
   [6, 6, 6]]]]

OutputIndexTensor: (Sizes:{1,1,3,3}, DataType:UINT32)
[[[[3, 1, 2],
   [2, 3, 1],
   [0, 1, 2]]]]

Remarks

A newer version of this operator, DML_TOP_K1_OPERATOR_DESC, was introduced in DML_FEATURE_LEVEL_2_1.

Availability

This operator was introduced in DML_FEATURE_LEVEL_2_0.

Tensor constraints

  • InputTensor, OutputIndexTensor, and OutputValueTensor must have the same DimensionCount.
  • InputTensor and OutputValueTensor must have the same DataType.

Tensor support

DML_FEATURE_LEVEL_5_0 and above

Tensor Kind Supported dimension counts Supported data types
InputTensor Input 1 to 8 FLOAT32, FLOAT16, INT64, INT32, INT16, INT8, UINT64, UINT32, UINT16, UINT8
OutputValueTensor Output 1 to 8 FLOAT32, FLOAT16, INT64, INT32, INT16, INT8, UINT64, UINT32, UINT16, UINT8
OutputIndexTensor Output 1 to 8 UINT64, UINT32

DML_FEATURE_LEVEL_3_1 and above

Tensor Kind Supported dimension counts Supported data types
InputTensor Input 1 to 8 FLOAT32, FLOAT16, INT32, INT16, INT8, UINT32, UINT16, UINT8
OutputValueTensor Output 1 to 8 FLOAT32, FLOAT16, INT32, INT16, INT8, UINT32, UINT16, UINT8
OutputIndexTensor Output 1 to 8 UINT32

DML_FEATURE_LEVEL_2_1 and above

Tensor Kind Supported dimension counts Supported data types
InputTensor Input 4 FLOAT32, FLOAT16, INT32, INT16, INT8, UINT32, UINT16, UINT8
OutputValueTensor Output 4 FLOAT32, FLOAT16, INT32, INT16, INT8, UINT32, UINT16, UINT8
OutputIndexTensor Output 4 UINT32

DML_FEATURE_LEVEL_2_0 and above

Tensor Kind Supported dimension counts Supported data types
InputTensor Input 4 FLOAT32, FLOAT16
OutputValueTensor Output 4 FLOAT32, FLOAT16
OutputIndexTensor Output 4 UINT32

Requirements

Requirement Value
Header directml.h