Implement argMax and argMin #62

BruceDai · 2023-12-12T07:00:14Z

@huningxin @fdwr @miaobin PTAL, thanks.

src/arg_max_min.js

…nd reduction ops

test/reduce_test.js

test/arg_max_min_test.js

src/arg_max_min.js

BruceDai · 2023-12-13T07:09:24Z

@huningxin I've updated commit to address your comments, please take another look, thanks.

huningxin

lgtm

fdwr

Some corner case considerations and minor suggestions, else LGTM.

src/arg_max_min.js

fdwr · 2023-12-15T05:01:11Z

src/arg_max_min.js

+    } = {}) {
+  // If axes doesn't present (defaulting to null), all dimensions are reduced.
+  // See https://webmachinelearning.github.io/webnn/#dom-mlargminmaxoptions-axes.
+  const inpAxes = axes ?? new Array(input.rank).fill(0).map((_, i) => i);


Regarding axes ??, I'm not that familiar with the falsesyness of JS evaluation. If you pass axes = [], will that correctly use the explicit empty axes (output shape is then the same size as the input, and there is a 1:1 mapping of output indices) rather than generate full reduction?

axes = []

It's so very special, it's neither null nor a sequence having values in the range [0, N-1]. Is there any model using this axes?

According to https://webmachinelearning.github.io/webnn/#dom-mlargminmaxoptions-axes:

axes, of type sequence<unsigned long>, defaulting to null A sequence of [unsigned long](https://webidl.spec.whatwg.org/#idl-unsigned-long). The dimensions to reduce. The values in the sequence must be in the range [0, N-1] where N is the [rank](https://webmachinelearning.github.io/webnn/#mloperand-rank) of the input tensor. If not present, all dimensions are reduced.

My argument is mainly driven by robustness plus consistency with similar behavior for the reduction operators when passed explicit [] axes. You probably won't find a model unless you look really hard, but some model generation parameters could programmatically produce it; and we see that in related reduction operators like ReduceSum that explicit empty axes [] was important enough that ONNX specifically added the noop_with_empty_axes attribute later https://onnx.ai/onnx/operators/onnx__ReduceSum.html. Speaking of that, do we have any explicit axes = [] cases in reduceSum?

Last we also missed considering this axes = [] cases in reduceSum.
Since current Spec is lack of explanation for axes = [], I will open an issue for discussing it on WG.

I will open an issue for discussing it on WG.

Alrighty, and I will add my clear expectation of what should happen in the nop case 😉 (especially if WebNN is to support existing frameworks behavior where [] means nop).

Link to webmachinelearning/webnn#493

That was fast! I tried ONNX and TF:

x = tf.constant([[0,100],[200,300]], dtype=tf.float32) y = tf.math.reduce_sum(x, axis=[]) # https://www.tensorflow.org/api_docs/python/tf/math/reduce_sum print("value:", y) print("shape:", y.shape) # Prints: # value: tf.Tensor( # [[ 0. 100.] # [200. 300.]], shape=(2, 2), dtype=float32) # shape: (2, 2)

(will add that to the issue...)

Do you want to wait for the spec discussion or proceed and react later accordingly? (I don't mind either way, and you might not get a reply for the next 2 weeks...)

Thanks. I will follow your guide to support axes=[] here.

p.s. Appears that the Chromium implementation supports explicit axes = [] gracefully:
https://chromium-review.googlesource.com/c/chromium/src/+/5125652

src/reduce.js

fdwr · 2023-12-15T05:08:56Z

src/reduce.js

+
+  const valuesToReduce = [];
+  // Find all values to reduce.
+  for (let reduceIndex = 0; reduceIndex < reduceElements; ++reduceIndex) {


How about reducedElementCount? reduceElements sounds like it's actually the elements themselves (the values).

fdwr · 2023-12-15T05:16:00Z

src/reduce.js

+  const reduceDims = axes.map((axis) => input.shape[axis]);
+  const reduceElements = sizeOfShape(reduceDims);
+  const reduceStrides = new Array(axes.length);
+  reduceStrides[reduceStrides.length - 1] = 1;


Appears this will fail on the axes = [] case, which should return exactly one element of shape []. Most of the logic already looks correct to handle that case (including sizeOfShape which will return 1).

It's strange. It can work for this case without changes, my understanding is that index should be start from 0.

reduceStrides = []; reduceStrides[-1] = 1; // it can work

But I add if condition block before this, then it never goes here for axes=[], please see the update commit, thanks.

It's strange. It can work for this case without changes
reduceStrides[-1] = 1; // it can work

(huh, what a weird language: https://stackoverflow.com/questions/58570277/adding-an-element-in-negative-index-of-an-array)

reduceStrides[-1] = 1 getting [-1: 1] is bad, we should avoid that.

Thanks @huningxin!
I've updated commit to avoid it, please take another look, thanks.

fdwr · 2023-12-15T05:22:05Z

test/arg_max_min_test.js

+        },
+        argMax,
+        {
+          axes: [1, 0],


👍 Good that you have an axis-reversed case too in the argMinMax tests.

fdwr · 2023-12-15T05:30:46Z

test/arg_max_min_test.js

+          selectLastIndex: true,
+        });
+  });
+});


For completeness, we should include an empty axes case:

it('argMin axes=[]', function() { testArgMaxMin( { shape: [3, 3], value: [ 1, 2, 3, 3, 0, 0, 2, 5, 2, ], }, { shape: [3, 3], value: [ 0, 0, 0, 0, 0, 0, 0, 0, 0, ], }, argMin, { axes: [], }); });

And the corner case for scalar (which can only ever return a scalar of 0):

it('argMin scalar axes=[]', function() { testArgMaxMin( { shape: [], value: [3], }, { shape: [], value: [0], }, argMin, { axes: [], }); });

(note in these cases, keepDimensions and selectLastIndex have no effect on the output)

BruceDai · 2023-12-16T01:52:42Z

@fdwr I've updated commit to address your comments, please take another look. Thanks a lot!

fdwr

LGTM. Thanks Bruce. ⭐

fdwr · 2023-12-16T04:14:15Z

src/reduce.js

+  const reduceDims = axes.map((axis) => input.shape[axis]);
+  const reduceElements = sizeOfShape(reduceDims);
+  const reduceStrides = new Array(axes.length);
+  reduceStrides[reduceStrides.length - 1] = 1;


It's strange. It can work for this case without changes
reduceStrides[-1] = 1; // it can work

(huh, what a weird language: https://stackoverflow.com/questions/58570277/adding-an-element-in-negative-index-of-an-array)

fdwr · 2023-12-16T04:22:51Z

test/arg_max_min_test.js

+        });
+  });
+
+  it('argMax scalar axes=[] none effect by both keepDimensions and selectLastIndex being true',


no effect

huningxin · 2023-12-16T08:57:12Z

src/reduce.js

+  const reduceDims = axes.map((axis) => input.shape[axis]);
+  const reducedElementCount = sizeOfShape(reduceDims);
+  const reduceStrides = new Array(axes.length);
+  reduceStrides[reduceStrides.length - 1] = 1;


Suggested change

reduceStrides[reduceStrides.length - 1] = 1;

if (reduceStrides.length > 0) {

reduceStrides[reduceStrides.length - 1] = 1;

}

fdwr

LGTM with Ningxin's comment applied.

Implement argMax and argMin

9856a81

miaobin reviewed Dec 12, 2023

View reviewed changes

src/arg_max_min.js Outdated Show resolved Hide resolved

Fix output shape issues when keepDimensions=false for argMax/argMin a…

cf4cae7

…nd reduction ops

BruceDai mentioned this pull request Dec 13, 2023

Implement webnn-baseline and WPT for operations needed for well-known transformers webmachinelearning/webnn#485

Closed

huningxin reviewed Dec 13, 2023

View reviewed changes

test/reduce_test.js Outdated Show resolved Hide resolved

test/arg_max_min_test.js Show resolved Hide resolved

src/arg_max_min.js Outdated Show resolved Hide resolved

Address comments of adding more tests and enhancing readability

ff919e5

huningxin approved these changes Dec 13, 2023

View reviewed changes

fdwr reviewed Dec 15, 2023

View reviewed changes

BruceDai mentioned this pull request Dec 15, 2023

Clarify the expected behavior for axes being empty [] webmachinelearning/webnn#493

Closed

Address comments of supporting axes=[] and code optimization

857871f

BruceDai force-pushed the add_argmax_argmin branch from da3fcce to 857871f Compare December 16, 2023 02:19

fdwr approved these changes Dec 16, 2023

View reviewed changes

fix JS sort() by lambda

9321913

huningxin reviewed Dec 16, 2023

View reviewed changes

fdwr approved these changes Dec 16, 2023

View reviewed changes

Avoid adding an element in negative index of an array

7fe0842

fdwr approved these changes Dec 17, 2023

View reviewed changes

huningxin approved these changes Dec 18, 2023

View reviewed changes

huningxin merged commit 058aa77 into webmachinelearning:main Dec 18, 2023
3 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement argMax and argMin #62

Implement argMax and argMin #62

BruceDai commented Dec 12, 2023

BruceDai commented Dec 13, 2023

huningxin left a comment

fdwr left a comment

fdwr Dec 15, 2023

BruceDai Dec 15, 2023

fdwr Dec 15, 2023 •

edited

Loading

BruceDai Dec 15, 2023

fdwr Dec 15, 2023

BruceDai Dec 15, 2023

fdwr Dec 15, 2023

fdwr Dec 15, 2023

BruceDai Dec 15, 2023

fdwr Dec 15, 2023

fdwr Dec 15, 2023

BruceDai Dec 16, 2023

fdwr Dec 15, 2023

BruceDai Dec 16, 2023

fdwr Dec 16, 2023 •

edited

Loading

huningxin Dec 16, 2023

BruceDai Dec 17, 2023

fdwr Dec 15, 2023

huningxin Dec 16, 2023

fdwr Dec 15, 2023

BruceDai commented Dec 16, 2023

fdwr left a comment

fdwr Dec 16, 2023 •

edited

Loading

fdwr Dec 16, 2023

huningxin Dec 16, 2023

fdwr left a comment

Implement argMax and argMin #62

Implement argMax and argMin #62

Conversation

BruceDai commented Dec 12, 2023

BruceDai commented Dec 13, 2023

huningxin left a comment

Choose a reason for hiding this comment

fdwr left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fdwr Dec 15, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fdwr Dec 16, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

BruceDai commented Dec 16, 2023

fdwr left a comment

Choose a reason for hiding this comment

fdwr Dec 16, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fdwr left a comment

Choose a reason for hiding this comment

fdwr Dec 15, 2023 •

edited

Loading

fdwr Dec 16, 2023 •

edited

Loading

fdwr Dec 16, 2023 •

edited

Loading