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. 😉

... 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.

[huningxin]

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

[huningxin]

The axis needs to be validated according to input's rank.

[inexorabletash]

Added for the operator version in 99e8773

For the activation version - would that be done synchronously when the activation is provided in a graph builder call, or at build time? Do we have other cases of activation validation?

[fdwr]

For the activation version - would that be done synchronously when the activation is provided in a graph builder call,

All the information should be available at construction time (pre-build), since WebNN requires explicit shapes during construction.

Do we have other cases of activation validation?

None come to mind, except possibly softplus steepness, but that's not impacted by the tensor shape constraints.

[huningxin]

For the activation version - would that be done synchronously when the activation is provided in a graph builder call,

All the information should be available at construction time (pre-build), since WebNN requires explicit shapes during construction.

Agreed. In particular, I suppose the validation should be done at the build method of operators that accept a fused activation function, like conv2d.

[inexorabletash]

Here's a sketch of what this could look like:

• In create an MLActivation, add optional validation steps parameter, store the validation steps in an internal slot. The default validation steps are to return true.
• When creating softmax activation, pass these steps as the validation steps:
  1. If axis is greater than or equal to input’s rank, then return false.
  2. Otherwise, return true.
  • Note that if this were C++, this would be a lambda function, capturing axis.
• In all 7 methods that take activation functions (singular or multiple), add steps like this:
  1. If options.activation exists, and running its validation steps with input returns false, then throw a TypeError.

Thoughts:

• This could also be bundled into the validate activation steps - make those take an MLOperand in addition to or instead of MLGraphBuilder.
• Is passing a single MLOperand (i.e. input) sufficient for validation?

WDYT?

[huningxin]

• Is passing a single MLOperand (i.e. input) sufficient for validation?

I suppose it should pass output MLOperand since the activation follows the operator that fuses it.

[inexorabletash]

I guess we'll need to pass the output descriptor.

[fdwr]

In particular, I suppose the validation should be done at the build method of operators that accept a fused activation function, like conv2d.

Yep, since stand-alone activations adopt whatever the size is of the operator they are joined to or used inside. So full validation must be deferred for these.

MLConv2dOptions::activation
MLConvTranspose2dOptions::activation
MLBatchNormalizationOptions::activation
MLGruOptions::activations
MLGruCellOptions::activations
MLLstmOptions::activations
MLLstmCellOptions::activations

Is passing a single MLOperand (i.e. input) sufficient for validation?

Sounds fine, the MLActivation and at least either the MLOperand.shape() or the MLOperand itself for more generality.

[inexorabletash]

Okay - 7981fa5 sketches out the validation. Notes:

• I went with passing an MLOperandDescriptor, because earlier we decided that we'd do all validation before creating the operator and output MLOperand to pass.
• The phrasing around creating/passing a lambda is based on Fetch
• The GRU/LSTM outputs are complicated. I'm not sure I'm passing the right thing, so please please double check.

[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 66bf64a

[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.

[fdwr]

3 suggestions, and rest looks good to me. Thanks Josh.

[fdwr]

LGTM given Ningxin's validation comment and some single-click ```suggestion blocks.

[fdwr]

LGTM. TYJB. 👌

[inexorabletash]

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

[inexorabletash]

• Converting this to a draft since gru/lstm: Fix missing/incorrect variables and options use #655 should land first.

[fdwr]

👍

[inexorabletash]

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

[huningxin]

LGTM

[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 -_-).