Modularize argMax, argMin, min, mean and sum #3517

tafsiri · 2020-06-26T03:53:02Z

Also move norm and max tests into their own files.

To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.

This change is

tafsiri · 2020-06-29T15:46:07Z

Modularizes the rest of the reduction ops and makes the max kernels handle the keepDims attr across backends.

It also fixes a bug that was present in the max kernel in wasm where if an internal transpose was used the input tensor would be changed after the op completed. I added a test for that to max and min, but a general thing to watch out for is calling .set on the typedarray for a models inputs.

annxingyuan

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

tfjs-core/src/kernel_names.ts, line 386 at r4 (raw file):

export const Mean = 'Mean';
export type MeanInputs = Pick<NamedTensorInfoMap, 'x'>;
export interface MeanAttrs {

Maybe would make sense to have ReductionInputs and ReductionAttrs (similar to BinaryInputs) given that they share the same interface.

tfjs-core/src/ops/arg_max.ts, line 62 at r4 (raw file):

    const permutedAxes = axis_util.getAxesPermutation(axes, $x.rank);
    if (permutedAxes != null) {
      $x = $x.transpose(permutedAxes);

remove chained op

tfjs-core/src/ops/arg_max.ts, line 66 at r4 (raw file):

    }
    const res = backend.argMax($x, axes[0]);
    return res;

nit: return backend.argMax

tfjs-core/src/ops/arg_min.ts, line 66 at r4 (raw file):

    const permutedAxes = axis_util.getAxesPermutation(axes, $x.rank);
    if (permutedAxes != null) {
      $x = $x.transpose(permutedAxes);

remove chained op

tfjs-core/src/ops/mean.ts, line 80 at r4 (raw file):

    const value = sum(res, axis, keepDims);

    const gradFunc = (dy: Tensor) => {

Would we want to eventually modularize custom grad functions as well?

tfjs-core/src/ops/sum.ts, line 62 at r4 (raw file):

  let $x = convertToTensor(x, 'x', 'sum');
  if ($x.dtype === 'bool') {
    $x = $x.toInt();

Sorry if we already discussed this - did we consider whether casting should happen at the kernel level?

tfjs-core/src/ops/sum.ts, line 73 at r4 (raw file):

    let permutedX = $x;
    if (permutation != null) {
      permutedX = $x.transpose(permutation);

remove chained op

tfjs-core/src/ops/sum.ts, line 79 at r4 (raw file):

    if (keepDims) {
      const newShape = expandShapeToKeepDim(value.shape, axes);
      value = value.reshape(newShape);

remove chained

tfjs-core/src/public/chained_ops/arg_max.ts, line 29 at r4 (raw file):

Tensor.prototype.argMax = function<T extends Tensor>(axis?: number): T {
  this.throwIfDisposed();
  // tslint:disable-next-line: no-unnecessary-type-assertion

could you include a comment about why the assertion is nevertheless needed? curious why min seems to be the exception in that it doesn't need an assertion?

tafsiri

Reviewed 52 of 52 files at r1, 3 of 3 files at r2, 17 of 17 files at r3.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan and @lina128)

tfjs-core/src/kernel_names.ts, line 386 at r4 (raw file):