Fix argMax & argMin benchmarks

[davidsoergel]

Description

The reduction op benchmarks operate on a 2D tensor and expect a scalar response. However, argMax and argMIn reduce rank by only one dimension at a time. This fix simply applies these ops twice in a row in order to produce a scalar.

It may seem a little weird that we'll now measure time for two calls instead of one, but I think that's OK.

• it doesn't matter as long as we consider this a relative measure,
• this makes the measurement more comparable to the other reduction ops, which are reducing over both dimensions here, and
• the time is likely dominated by the first call anyway.

Fixes tensorflow/tfjs#504

BUG

---

For repository owners only:

Please remember to apply all applicable tags to your pull request.
Tags: FEATURE, BREAKING, BUG, PERF, DEV, DOC, SECURITY

For more info see: https://github.com/tensorflow/tfjs/blob/master/DEVELOPMENT.md

---

This change is 

[nsthorat]

Hmm -- just curious, what is the motivation for this change?

[davidsoergel]

The motivation came from the linked bug. Running the benchmark on argMin or argMax currently crashes with Error: Number of coordinates in get() must match the rank of the tensor, because reduction_ops_benchmark:50 has op(input).get();, which only works if op(input) is a scalar.

Sure we could split the test into two (one for to-scalar reducers and one for to-rank-minus-one reducers), but this solution seemed more expedient.

[nsthorat]

How about instead we just call ".data()" for reduction ops?

[nsthorat]

Or, actually even better, call reduction ops on a vector instead of matrix?

[nsthorat]

Gentle ping on this since I am also now seeing this.

[davidsoergel]

Oops, sorry I lost track of this. Implemented the 1D approach.

[dsmilkov]

Reviewed 2 of 2 files at r2.
Reviewable status: 0 of 1 approvals obtained (waiting on @nsthorat and @dsmilkov)