Add Kernel SparseReshape for CPU and WebGL backend #4956

pyu10055 · 2021-04-19T22:39:05Z

ref #4838
TensorFlow python version is defined here
c++ kernel definition is here
To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.

This change is

lina128

Hi Ping, I left some minor comments, overall LGTM. Will converter code be added in a separate PR?

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @pyu10055)

tfjs-backend-cpu/src/kernels/SparseReshape.ts, line 29 at r1 (raw file):

inputIndices.shape.length <= 1

The equivalent of isMatrix is if(inputIndices.shape.length === 2)

tfjs-backend-cpu/src/kernels/SparseReshape.ts, line 50 at r1 (raw file):

      Array.from(backend.data.get(newShape.dataId).values as TypedArray);

  const [newValues, indicesShape, outputShape] = sparseReshapeImpl(

nit: rename to newIndices

tfjs-backend-cpu/src/kernels/SparseReshape_impl.ts, line 57 at r1 (raw file):

          'specified input sizes are non-zero');
    }
    const missing = Math.round(denseSize / product);

Should it be Math.trunc?

tfjs-backend-cpu/src/kernels/SparseReshape_impl.ts, line 90 at r1 (raw file):

  }

  const newValues =

nit: rename to newIndices

tfjs-backend-cpu/src/kernels/SparseReshape_impl.ts, line 92 at r1 (raw file):

  const newValues =
      util.getArrayFromDType(inputDType, nnz * outputRank) as TypedArray;
  util.getArrayFromDType(inputDType, nnz * outputRank);

Remove this line.

tfjs-backend-cpu/src/kernels/SparseReshape_impl.ts, line 96 at r1 (raw file):

i * inputRank + j

Is it more readable to write i * inputIndicesStrides[0] + j, where const inputIndicesStrides = util.computeStrides(inputIndicesShape);
Alternatively, add annotation that inputIndices is a matrix, which has shape [numOfIndices, rank].

tfjs-backend-cpu/src/kernels/SparseReshape_impl.ts, line 99 at r1 (raw file):

Math.floor(id / outputStrides[j]);

nit: Math.trunc

tfjs-backend-webgl/src/kernels/SparseReshape.ts, line 28 at r1 (raw file):

  const {inputs, backend} = args;
  const {inputIndices, inputShape, newShape} = inputs;
  if (inputIndices.shape.length <= 1) {

!== 2

tfjs-backend-webgl/src/kernels/SparseReshape.ts, line 48 at r1 (raw file):

newValues

nit: rename to newIndices

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

}

export const SparseReshape = 'SparseReshape';

Try to order alphabetically, add this definition after SpaceToBatchND.

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

Quoted 5 lines of code…

export interface SparseReshapeInputs extends NamedTensorInfoMap {
  inputIndices: TensorInfo;
  inputShape: TensorInfo;
  newShape: TensorInfo;
}

Why not just export type SparseReshapeInputs = Pick<NamedTensorInfoMap, 'inputIndices'|'inputShape'|'newShape'>;

tfjs-core/src/ops/ops.ts, line 228 at r1 (raw file):

import * as fused from './fused_ops';
import * as sparse from './sparse_ops';

Can we follow the Image Ops import and export below? Because we don't want to export everything from ./sparse_ops, e.g. impl shouldn't be exported.

tfjs-core/src/ops/sparse_ops.ts, line 18 at r1 (raw file):

 */

import {sparseReshape} from './sparse/sparse_reshape';

This feels redundant, should we just manage all the import export in ops.ts?

pyu10055

Thank you, will add converter changes in a separate PR.

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @lina128)

tfjs-backend-cpu/src/kernels/SparseReshape.ts, line 29 at r1 (raw file):