Repository: microsoft/onnxruntime
File: onnxruntime/core/providers/webnn/builders/impl/softmax_op_builder.cc (C++)

Code before:
  emscripten::val input = model_builder.GetOperand(node.InputDefs()[0]->Name());
  std::vector<int64_t> input_shape;
  ORT_RETURN_IF_NOT(GetShape(*input_defs[0], input_shape, logger), "Cannot get shape");
  const auto input_size = input_shape.size();

  emscripten::val options = emscripten::val::object();

  NodeAttrHelper helper(node);
  const auto since_version = node.SinceVersion();
  const int32_t default_axis = since_version < 13 ? 1 : -1;
  int32_t axis = helper.Get("axis", default_axis);
  axis = SafeInt<int32_t>(HandleNegativeAxis(axis, input_size));

  // Prior to opset 13, Softmax operates with different semantics compared to opset 13 and later.
  // Specifically, it normalizes over the flattened range of dimensions starting from the specified
  // axis to the last dimension.
  // In contrast, WebNN's softmax aligns with the behavior introduced in opset 13 and later.
  // To handle the differences for earlier opsets, a reshape operation can be applied if necessary.
  const bool do_reshape = since_version < 13 && axis != SafeInt<int32_t>(input_size - 1);
  std::vector<uint32_t> input_shape_uint32;
  if (do_reshape) {
    input_shape_uint32 = GetNarrowedIntFromInt64<uint32_t>(input_shape);
    // Need to reshape the input to 2D tensor with new shape [M, N].
    // M = d0*d1*...*d(axis-1), N = d(axis)*...*d(n-1)
    const auto M = axis == 0
                       ? 1
                       : Product(std::vector<uint32_t>(input_shape_uint32.begin(), input_shape_uint32.begin() + axis));
    const auto N = Product(std::vector<uint32_t>(input_shape_uint32.begin() + axis, input_shape_uint32.end()));
    emscripten::val new_shape = emscripten::val::array();
    new_shape.set(0, M);
    new_shape.set(1, N);

    options.set("label", node.Name() + "_reshape_input");
    input = model_builder.GetBuilder().call<emscripten::val>("reshape", input, new_shape, options);
    // Apply softmax along the last dimension (N).
    axis = 1;
  }

  options.set("label", node.Name());
  emscripten::val output = model_builder.GetBuilder().call<emscripten::val>("softmax", input, axis, options);

  if (do_reshape) {
    // Softmax has the same output shape as input shape.
    // Reshape the output back to the original input shape.
    options.set("label", node.Name() + "_reshape_output");
    output = model_builder.GetBuilder().call<emscripten::val>(
        "reshape", output, emscripten::val::array(input_shape_uint32), options);
  }

  model_builder.AddOperand(node.OutputDefs()[0]->Name(), std::move(output));
  return Status::OK();

Review comment (suggestion):
Does `axis == 0` need a special case? Isn't `product(emptyArray)` == 1 anyway?

https://github.com/microsoft/onnxruntime/blob/f7fd3b522d100482988194f3ecc905102b83032f/onnxruntime/core/providers/common.h#L179-L181

Code after:
  emscripten::val input = model_builder.GetOperand(node.InputDefs()[0]->Name());
  std::vector<int64_t> input_shape;
  ORT_RETURN_IF_NOT(GetShape(*input_defs[0], input_shape, logger), "Cannot get shape");
  const auto input_size = input_shape.size();

  emscripten::val options = emscripten::val::object();

  NodeAttrHelper helper(node);
  const auto since_version = node.SinceVersion();
  const int32_t default_axis = since_version < 13 ? 1 : -1;
  int32_t axis = helper.Get("axis", default_axis);
  axis = SafeInt<int32_t>(HandleNegativeAxis(axis, input_size));

  // Prior to opset 13, Softmax operates with different semantics compared to opset 13 and later.
  // Specifically, it normalizes over the flattened range of dimensions starting from the specified
  // axis to the last dimension.
  // In contrast, WebNN's softmax aligns with the behavior introduced in opset 13 and later.
  // To handle the differences for earlier opsets, a reshape operation can be applied if necessary.
  const bool do_reshape = since_version < 13 && axis != SafeInt<int32_t>(input_size - 1);
  std::vector<uint32_t> input_shape_uint32;
  if (do_reshape) {
    input_shape_uint32 = GetNarrowedIntFromInt64<uint32_t>(input_shape);
    // Need to reshape the input to 2D tensor with new shape [M, N].
    // M = d0*d1*...*d(axis-1), N = d(axis)*...*d(n-1)
    const auto M = Product(std::vector<uint32_t>(input_shape_uint32.begin(), input_shape_uint32.begin() + axis));
    const auto N = Product(std::vector<uint32_t>(input_shape_uint32.begin() + axis, input_shape_uint32.end()));
    emscripten::val new_shape = emscripten::val::array();
    new_shape.set(0, M);
    new_shape.set(1, N);

    options.set("label", node.Name() + "_reshape_input");
    input = model_builder.GetBuilder().call<emscripten::val>("reshape", input, new_shape, options);
    // Apply softmax along the last dimension (N).
    axis = 1;
  }

  options.set("label", node.Name());
  emscripten::val output = model_builder.GetBuilder().call<emscripten::val>("softmax", input, axis, options);

  if (do_reshape) {
    // Softmax has the same output shape as input shape.
    // Reshape the output back to the original input shape.
    options.set("label", node.Name() + "_reshape_output");
    output = model_builder.GetBuilder().call<emscripten::val>(
        "reshape", output, emscripten::val::array(input_shape_uint32), options);
  }

  model_builder.AddOperand(node.OutputDefs()[0]->Name(), std::move(output));
  return Status::OK();
}