DML_ARGMAX_OPERATOR_DESC structure (directml.h)

Article
06/07/2022

Outputs the indices of the maximum-valued elements within one or more dimensions of the input tensor.

Each output element is the result of applying an argmax reduction on a subset of the input tensor. The argmax function outputs the index of the maximum-valued element within a set of input elements. The input elements involved in each reduction are determined by the provided input axes. Similarly, each output index is with respect to the provided input axes. If all input axes are specified, then the operator applies a single argmax reduction, and produces a single output element.

Syntax

struct DML_ARGMAX_OPERATOR_DESC {
  const DML_TENSOR_DESC *InputTensor;
  const DML_TENSOR_DESC *OutputTensor;
  UINT                  AxisCount;
  const UINT            *Axes;
  DML_AXIS_DIRECTION    AxisDirection;
};

Members

InputTensor

Type: const DML_TENSOR_DESC*

The tensor to read from.

OutputTensor

Type: const DML_TENSOR_DESC*

The tensor to write the results to. Each output element is the result of an argmax reduction on a subset of elements from the InputTensor.

DimensionCount must match InputTensor.DimensionCount (the rank of the input tensor is preserved).
Sizes must match InputTensor.Sizes, except for dimensions included in the reduced Axes, which must be size 1.

AxisCount

Type: UINT

The number of axes to reduce. This field determines the size of the Axes array.

Axes

Type: _Field_size_(AxisCount) const UINT*

The axes along which to reduce. Values must be in the range [0, InputTensor.DimensionCount - 1].

AxisDirection

Type: DML_AXIS_DIRECTION

Determines which index to select when multiple input elements have the same value.

DML_AXIS_DIRECTION_INCREASING returns the index of the first maximum-valued element (for example, argmax({3,2,1,2,3}) = 0)
DML_AXIS_DIRECTION_DECREASING returns the index of the last maximum-valued element (for example, argmax({3,2,1,2,3}) = 4)

Examples

The examples in this section all use this same two-dimensional input tensor.

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

Example 1. Applying argmax to columns

AxisCount: 1
Axes: {0}
AxisDirection: DML_AXIS_DIRECTION_INCREASING
OutputTensor: (Sizes:{1, 3}, DataType:UINT32)
[[1,  // argmax({1, 3, 2})
  2,  // argmax({2, 0, 5})
  1]] // argmax({3, 4, 2})

Example 2. Applying argmax to rows

AxisCount: 1
Axes: {1}
AxisDirection: DML_AXIS_DIRECTION_INCREASING
OutputTensor: (Sizes:{3, 1}, DataType:UINT32)
[[2], // argmax({1, 2, 3})
 [2], // argmax({3, 0, 4})
 [1]] // argmax({2, 5, 2})

Example 3. Applying argmax to all axes (the entire tensor)

AxisCount: 2
Axes: {0, 1}
AxisDirection: DML_AXIS_DIRECTION_INCREASING
OutputTensor: (Sizes:{1, 1}, DataType:UINT32)
[[7]]  // argmax({1, 2, 3, 3, 0, 4, 2, 5, 2})

Remarks

The output tensor sizes must be the same as the input tensor sizes, except for the reduced axes, which must be 1.

When AxisDirection is DML_AXIS_DIRECTION_INCREASING, this API is equivalent to DML_REDUCE_OPERATOR_DESC with DML_REDUCE_FUNCTION_ARGMAX.

A subset of this functionality is exposed through the DML_REDUCE_OPERATOR_DESC operator, and is supported on earlier DirectML feature levels.

Availability

This operator was introduced in DML_FEATURE_LEVEL_3_0.

Tensor constraints

InputTensor and OutputTensor must have the same DimensionCount.

Tensor support

DML_FEATURE_LEVEL_4_1 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
OutputTensor	Output	1 to 8	INT64, INT32, UINT64, UINT32

DML_FEATURE_LEVEL_3_0 and above

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

Requirements

Requirement	Value
Minimum supported client	Windows 10 Build 20348
Minimum supported server	Windows 10 Build 20348
Header	directml.h

Share via