Skip to content

Commit 4678f5d

Browse files
Honryrohan11235813
authored andcommitted
[WebNN EP] Add outputDataType option for the ArgMax/ArgMin ops (#21385)
### Description WebNN spec introduces a new option: `outputDataType` to `argMax` and `argMin` ops, it's default value is `int32`, we should explicitly set it to `int64` for WebNN EP. Spec CR: "Add outputDataType to argmin/argmax" webmachinelearning/webnn#730
1 parent 8c7423a commit 4678f5d

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

onnxruntime/core/providers/webnn/builders/impl/argmax_min_op_builder.cc

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,11 @@ Status ArgMaxMinOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
5050
options.set("axes", axes);
5151
options.set("keepDimensions", keep_dims == 1);
5252
options.set("selectLastIndex", select_last_index == 1);
53+
// TODO: use WebNN's opSupportLimits API to check the backend's supported output data types.
54+
// If the backend doesn't support int64 output, we should use default int32 output data type
55+
// then do a type casting (int32 -> int64) for the output. Refer to the CoreML EP for how to
56+
// support int64 output.
57+
options.set("outputDataType", "int64");
5358
emscripten::val output = emscripten::val::object();
5459

5560
const auto& op_type = node.OpType();

0 commit comments

Comments
 (0)