Skip to content

Commit

Permalink
[WebNN EP] Add outputDataType option for the ArgMax/ArgMin ops (#21385)
Browse files Browse the repository at this point in the history
### 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
  • Loading branch information
Honry authored Jul 22, 2024
1 parent a6c5e2c commit 17e9ea6
Showing 1 changed file with 5 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,11 @@ Status ArgMaxMinOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
options.set("axes", axes);
options.set("keepDimensions", keep_dims == 1);
options.set("selectLastIndex", select_last_index == 1);
// TODO: use WebNN's opSupportLimits API to check the backend's supported output data types.
// If the backend doesn't support int64 output, we should use default int32 output data type
// then do a type casting (int32 -> int64) for the output. Refer to the CoreML EP for how to
// support int64 output.
options.set("outputDataType", "int64");
emscripten::val output = emscripten::val::object();

const auto& op_type = node.OpType();
Expand Down

0 comments on commit 17e9ea6

Please sign in to comment.