Softmax axis absent

[fdwr]

(raised by @Honry in review https://github.com/microsoft/onnxruntime/pull/17665/files)

TF/PT/ONNX all take an axis parameter:

• TensorFlow tf.nn.softmax
• PyTorch torch.nn.Softmax
• ONNX Softmax-13

...but WebNN's softmax does not, making it challenging to implement a caller's softmax in terms of the function of the same name in WebNN. It is possible (see here) via bracketing transposes+reshapes around it, but the transpose+reshape contortions are unfortunate, and they could be more efficiently implemented in the backend rather than in each framework.

• ✅ Apple Metal Performance Shaders softMax has an axis.
• ✅ Apple MIL activation.softmax supports an axis.
• ✅ DirectML's DML_ACTIVATION_SOFTMAX1_OPERATOR_DESC supports an arbitrary axis list and dimensions, just like reduce. The older DML_ACTIVATION_SOFTMAX_OPERATOR_DESC can achieve it via reshapes/transpose/strides.
• ☑ XNNPack - limited to 2D input currently. This can be achieved by updating XNNPack to accept an axis or by using the existing XNNPack operator plus a reshape (in the simple case when the axis is the last dimension) or transpose (if the axis comes before the last dimension).

So it's achievable in each backend, even without any changes to the DML/XNNPack API's, but it would move the pain from the caller down to where it can be handled efficiently.

---

https://www.w3.org/TR/webnn/#api-mlgraphbuilder-softmax-method

partial interface MLGraphBuilder {
-  MLOperand softmax(MLOperand input);
+  MLOperand softmax(MLOperand input, unsigned long axis);
-  MLActivation softmax();
+  MLActivation softmax(unsigned long axis);
};

The behavior of this operation can be generically emulated from the usage of other operations as follow. However, user agents typically have a more efficient implementation for it, therefore its usage is encouraged from the performance standpoint.
// This sample deploys a well-known implementation trick [1] to compute the
// exponentials of the distances to the max value, instead of the exponentials
// of the input values itself, in order to increase the numerical stability of
// the result.
// [1]: https://cs231n.github.io/linear-classify/#softmax
- const maxX = builder.reduceMax(x, { axes: [1], keepDimensions: true });
+ const maxX = builder.reduceMax(x, { axes: [axis], keepDimensions: true });
const expX = builder.exp(builder.sub(x, maxX));
- return builder.div(expX, builder.reduceSum(expX, { axes: [1], keepDimensions: true }));
+ return builder.div(expX, builder.reduceSum(expX, { axes: [axis], keepDimensions: true }));

[huningxin]

Should it be:

-  MLOperand softmax(MLOperand input);
+  MLOperand softmax(MLOperand input, unsigned long axis);

[fdwr]

Ningxin: Indeed, I fixed my typo right before you wrote your comment 😅.