Modularize scatterNd, unsortedSegmentSum & gather.

tfjs-backend-wasm/src/kernels/Gather.ts → tfjs-backend-wasm/src/kernels/GatherV2.ts

@@ -15,19 +15,11 @@
  * =============================================================================
  */
 
-import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo, util} from '@tensorflow/tfjs-core';
+import {backend_util, GatherV2, GatherV2Attrs, GatherV2Inputs, KernelFunc, registerKernel, Tensor, TensorInfo, util} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
-import {CppDType} from './types';
-
-interface GatherInputs extends NamedTensorInfoMap {
-  x: TensorInfo;
-  indices: TensorInfo;
-}
 
-interface GatherAttrs extends NamedAttrMap {
-  axis: number;
-}
+import {CppDType} from './types';
 
 let wasmGather:
     (xId: number, dtype: CppDType, xStrides: Uint8Array, stridesSize: number,
@@ -47,8 +39,8 @@ function setup(backend: BackendWasm): void {
   ]);
 }
 
-function gather(
-    args: {backend: BackendWasm, inputs: GatherInputs, attrs: GatherAttrs}):
+function gatherV2(
+    args: {backend: BackendWasm, inputs: GatherV2Inputs, attrs: GatherV2Attrs}):
     TensorInfo {
   const {backend, inputs, attrs} = args;
   const {x, indices} = inputs;
@@ -80,12 +72,18 @@ function gather(
       xId, CppDType[x.dtype], xStridesBytes, stridesSize, indicesId, axis,
       outStridesBytes, outId);
 
+  // reshape
+  const parsedAxis = util.parseAxisParam(axis, x.shape)[0];
+  const shapeInfo = backend_util.segment_util.collectGatherOpShapeInfo(
+      x as Tensor, indices as Tensor, parsedAxis);
+
+  out.shape = shapeInfo.outputShape;
   return out;
 }
 
 registerKernel({
-  kernelName: 'Gather',
+  kernelName: GatherV2,
   backendName: 'wasm',
   setupFunc: setup,
-  kernelFunc: gather
+  kernelFunc: gatherV2 as {} as KernelFunc
 });

tfjs-backend-wasm/src/kernels/ScatterNd.ts

@@ -15,19 +15,11 @@
  * =============================================================================
  */
 
-import {NamedAttrMap, NamedTensorInfoMap, registerKernel, scatter_util, TensorInfo, util} from '@tensorflow/tfjs-core';
+import {KernelFunc, registerKernel, scatter_util, ScatterNd, ScatterNdAttrs, ScatterNdInputs, TensorInfo, util} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
-import {CppDType} from './types';
-
-interface ScatterNdInputs extends NamedTensorInfoMap {
-  indices: TensorInfo;
-  updates: TensorInfo;
-}
 
-interface ScatterNdAttrs extends NamedAttrMap {
-  shape: number[];
-}
+import {CppDType} from './types';
 
 let wasmScatterNd: (
     indicesId: number, updatesId: number, dtype: CppDType, sliceRank: number,
@@ -81,8 +73,8 @@ function scatterNd(
 }
 
 registerKernel({
-  kernelName: 'ScatterNd',
+  kernelName: ScatterNd,
   backendName: 'wasm',
   setupFunc: setup,
-  kernelFunc: scatterNd
+  kernelFunc: scatterNd as {} as KernelFunc
 });

tfjs-backend-wasm/src/kernels/all_kernels.ts

@@ -41,7 +41,7 @@ import './FloorDiv';
 import './FusedBatchNorm';
 import './FusedConv2D';
 import './FusedDepthwiseConv2D';
-import './Gather';
+import './GatherV2';
 import './GatherNd';
 import './Greater';
 import './GreaterEqual';

tfjs-core/src/gradients/GatherV2_grad.ts

@@ -0,0 +1,85 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {GatherV2, GatherV2Attrs} from '../kernel_names';
+import {GradConfig, NamedAttrMap} from '../kernel_registry';
+import {getUndoAxesPermutation} from '../ops/axis_util';
+import {reshape} from '../ops/reshape';
+import {transpose} from '../ops/transpose';
+import {unsortedSegmentSum} from '../ops/unsorted_segment_sum';
+import {Tensor, Tensor1D} from '../tensor';
+import {parseAxisParam} from '../util';
+
+export const gatherGradConfig: GradConfig = {
+  kernelName: GatherV2,
+  inputsToSave: ['x', 'indices'],
+  gradFunc: (dy: Tensor, saved: Tensor[], attrs: NamedAttrMap) => {
+    const [x, indices] = saved;
+    const {axis} = attrs as {} as GatherV2Attrs;
+
+    const parsedAxis = parseAxisParam(axis, x.shape)[0];
+
+    const derX = () => {
+      const paramsShape = x.shape;
+      const indicesSize = indices.size;
+
+      const outerShape = paramsShape.slice(0, parsedAxis);
+      const outerDims = outerShape.length;
+      const innerShape = paramsShape.slice(axis, paramsShape.length).slice(1);
+      const innerDims = innerShape.length;
+
+      const outerAxesIndices = arrayRange(0, outerDims);
+      const innerAxesIndices =
+          arrayRange(outerDims + 1, outerDims + 1 + innerDims);
+
+      const valuesShape = arrayConcat([outerShape, [indicesSize], innerShape]);
+
+      const values = reshape(dy, valuesShape);
+      const reshapedIndices = reshape(indices, [indicesSize]);
+
+      const transposeDims =
+          arrayConcat([[outerDims], outerAxesIndices, innerAxesIndices]);
+      const valuesTranspose = transpose(values, transposeDims);
+      let paramsGrad = unsortedSegmentSum(
+          valuesTranspose, reshapedIndices as Tensor1D, x.shape[parsedAxis]);
+
+      const invertTransposeDims = getUndoAxesPermutation(transposeDims);
+      paramsGrad = transpose(paramsGrad, invertTransposeDims);
+
+      return paramsGrad;
+    };
+    return {x: derX, indices: () => indices};
+  }
+};
+
+function arrayRange(start: number, stop: number): number[] {
+  const result = [];
+  for (let i = start; i < stop; ++i) {
+    result.push(i);
+  }
+  return result;
+}
+
+function arrayConcat(arrays: number[][]): number[] {
+  const result = [];
+  for (let i = 0; i < arrays.length; ++i) {
+    for (let j = 0; j < arrays[i].length; ++j) {
+      result.push(arrays[i][j]);
+    }
+  }
+  return result;
+}

tfjs-core/src/gradients/UnsortedSegmentSum_grad.ts

@@ -0,0 +1,56 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {UnsortedSegmentSum} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {expandDims} from '../ops/expand_dims';
+import {gather} from '../ops/gather';
+import {greaterEqual} from '../ops/greater_equal';
+import {logicalAnd} from '../ops/logical_and';
+import {maximum} from '../ops/maximum';
+import {ones, scalar, zerosLike} from '../ops/tensor_ops';
+import {where} from '../ops/where';
+import {Tensor, Tensor1D} from '../tensor';
+
+export const unsortedSegmentSumGradConfig: GradConfig = {
+  kernelName: UnsortedSegmentSum,
+  inputsToSave: ['segmentIds'],
+  gradFunc: (dy: Tensor, saved: Tensor[]) => {
+    const [segmentIds] = saved;
+
+    const derX = () => {
+      return gatherDropNegatives(dy, segmentIds as Tensor1D);
+    };
+    return {x: derX};
+  }
+};
+
+function gatherDropNegatives<T extends Tensor>(x: T, indices: Tensor1D) {
+  // Helper function for unsorted segment ops. Gathers params for
+  // positive segment ids and gathers 0 for inputs with negative segment id.
+  // Mirrors _GatherDropNegatives from tensorflow/python/ops/math_grad.py
+  const zeroClippedIndices = maximum(indices, zerosLike(indices));
+  const gathered = gather(x, zeroClippedIndices as Tensor1D);
+  let isPositive = greaterEqual(indices, scalar(0, 'int32'));
+  const numIters = gathered.rank - isPositive.rank;
+  for (let i = 0; i < numIters; ++i) {
+    isPositive = expandDims(isPositive, i + 1);
+  }
+  isPositive = logicalAnd(isPositive, ones(gathered.shape, 'bool'));
+  const zeroSlice = zerosLike(gathered);
+  return where(isPositive, gathered, zeroSlice);
+}

tfjs-core/src/kernel_names.ts

@@ -273,6 +273,12 @@ export interface FusedBatchNormAttrs {
   varianceEpsilon: number;
 }
 
+export const GatherV2 = 'GatherV2';
+export type GatherV2Inputs = Pick<NamedTensorInfoMap, 'x'|'indices'>;
+export interface GatherV2Attrs {
+  axis: number;
+}
+
 export const GatherNd = 'GatherNd';
 export type GatherNdInputs = Pick<NamedTensorInfoMap, 'params'|'indices'>;
 
@@ -499,6 +505,12 @@ export interface ReverseAttrs {
   dims: number|number[];
 }
 
+export const ScatterNd = 'ScatterNd';
+export type ScatterNdInputs = Pick<NamedTensorInfoMap, 'indices'|'updates'>;
+export interface ScatterNdAttrs {
+  shape: number[];
+}
+
 export const SelectV2 = 'SelectV2';
 export type SelectV2Inputs = Pick<NamedTensorInfoMap, 'condition'|'t'|'e'>;
 
@@ -553,6 +565,13 @@ export interface UnpackAttrs {
   axis: number;
 }
 
+export const UnsortedSegmentSum = 'UnsortedSegmentSum';
+export type UnsortedSegmentSumInputs =
+    Pick<NamedTensorInfoMap, 'x'|'segmentIds'>;
+export interface UnsortedSegmentSumAttrs {
+  numSegments: number;
+}
+
 /**
  * TensorFlow.js-only kernels
  */