DML_TOP_K1_OPERATOR_DESC structure (directml.h)

Selects the largest or smallest 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.

The choice of whether to select the largest K elements or the smallest K elements can be controlled using AxisDirection.

Syntax

struct DML_TOP_K1_OPERATOR_DESC {
  const DML_TENSOR_DESC *InputTensor;
  const DML_TENSOR_DESC *OutputValueTensor;
  const DML_TENSOR_DESC *OutputIndexTensor;
  UINT                  Axis;
  UINT                  K;
  DML_AXIS_DIRECTION    AxisDirection;
};

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. The top K elements are selected based on whether they are the largest or the smallest, depending on the value of AxisDirection. 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) if AxisDirection is DML_AXIS_DIRECTION_DECREASING. Otherwise, the opposite is true and the values selected are guaranteed to be sorted ascending (smallest to largest).

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. Note that this is true irrespective of the value of AxisDirection.

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.

AxisDirection

Type: DML_AXIS_DIRECTION

A value from the DML_AXIS_DIRECTION enumeration. If set to DML_AXIS_DIRECTION_INCREASING, then this operator returns the smallest K elements in order of increasing value. Otherwise, it returns the largest K elements in decreasing order.

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
AxisDirection: DML_AXIS_DIRECTION_DECREASING
   
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
AxisDirection: DML_AXIS_DIRECTION_DECREASING
   
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
AxisDirection: DML_AXIS_DIRECTION_DECREASING
   
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]]]]

Example 4. Increasing axis direction

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

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

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

Remarks

When AxisDirection is set to DML_AXIS_DIRECTION_DECREASING, this operator is equivalent to DML_TOP_K_OPERATOR_DESC.

Availability

This operator was introduced in DML_FEATURE_LEVEL_2_1.

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

DML_FEATURE_LEVEL_3_1 and above

DML_FEATURE_LEVEL_2_1 and above

Requirements