[WebNN EP] Add outputDataType option for the ArgMax/ArgMin ops

microsoft · Jul 19, 2024 · ca89be0 · ca89be0
1 parent 9140d9b
commit ca89be0
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/onnxruntime/core/providers/webnn/builders/impl/argmax_min_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/argmax_min_op_builder.cc
@@ -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();