Add axis argument to softmax()

[inexorabletash]

Frameworks (TensorFlow, PyTorch, ONNX) all accept an axis parameter.

Most backends also support an axis, or it can be emulated with a reshape. As @fdwr wrote: So it's achievable in each backend... but it would move the pain from the caller down to where it can be handled efficiently.

Fixes #466

---

Preview | Diff

[fdwr]

Thank you for championing my own issue in :).

FYI @Honry as this impacts your softmax implementation.

[fdwr]

A tangential high-level consideration (cc @anssiko too) - as we make changes to existing operators, the divergence between the Chromium implementation and the spec will bite people. I wonder how we point people to specific frozen versions of the spec to say Chrome/Edge v3.14159 matches with spec v42? Otherwise they'll try softmax from WebNN with an axis parameter, and it fails. The spec is still a "living spec" candidate recommendation, not frozen, but it becomes increasingly challenging the older the operator is (the fillSequence operator is less impactful because it didn't exist yet anyway) and so clear messaging helps users.

It appears we have versions produced near daily (or as often as changes are made) like here https://www.w3.org/TR/2024/CRD-webnn-20240416/ with an alarmingly red disclaimer that the spec is red. So maybe we can point to those for a mapping. Are they persistent for some time frame?

[inexorabletash]

Thank you for championing my own issue in :).

Just trying to ~watch the world burn~ burn the issue list down. :wink:

... the divergence between the Chromium implementation and the spec will bite people.

See #650 for a more severe case - it's going to be painful. That's one of the reasons I'm trying to pick off lots of issues - get the pain behind us, instead of in front of us. Tracking Canary vs. spec dates is probably too much work, but we could try and match the big releases (e.g. Chrome/Edge 124 mostly matches spec v20240101) ? I don't know if we can do that historically, but maybe that's work we take on as part of upcoming browser branch dates?

[huningxin]

this impacts your softmax implementation.

The ONNXRuntime WebNN EP softmax implementation also does the reshape N-D input to 2-D and reshape the 2-D output to N-D. Should WebNN softmax supports N-D input? It looks like DirectML and CoreML support N-D input.

[inexorabletash]

I cannot add inline comment, but "#### {{MLGraphBuilder/softmax(input)}} #### {#api-mlgraphbuilder-softmax-input}" at line 5300 needs to be updated with axis parameter.

Done in 66bf64a4ce6305cffb2920be626fdd8d4f22b273

[fdwr]

@huningxin: Does CoreML? I'm only seeing a classcoremltools.converters.mil.mil.ops.defs.iOS15.activation.softmax with an int32 axis. DirectML does, as axes are a more general superset of the single axis case, and since softmax is a div and reduceSumExp making it consistent with other reduction operators. However, while I like the elegance of axes, no other front-end or back-end I know of does so. So that's why I proposed a single axis.

I misread 😑 - yes, the input should be ND - I'll add a comment to that line.

[huningxin]

Does CoreML? I'm only seeing a classcoremltools.converters.mil.mil.ops.defs.iOS15.activation.softmax with an int32 axis.

@fdwr , I meant whether the input should be N-D. I am fine with the single axis proposal.

[inexorabletash]

I noticed some related glitches when putting 622ce7535c9a4ded50526a12e905b5bf75b36ecc together - I'll factor those out into a separate PR

[inexorabletash]

• Converting this to a draft since #655 should land first.

[inexorabletash]

Okay, back to a real PR. Final review @huningxin ? Do we want to discuss this more in the WG before merging?

[huningxin]

@inexorabletash

Okay, back to a real PR. Final review @huningxin ?

Thanks! Done.

Do we want to discuss this more in the WG before merging?

I think we can merge it. Any thoughts @fdwr ?

[fdwr]

I think we can merge it. Any thoughts @fdwr ?

@huningxin What are we discussing? It already looked good to me 👍. So I'll merge it now (and then sleep -_-).