tensorflow · pyu10055 · Jul 8, 2021 · Jul 3, 2021 · Jul 8, 2021
diff --git a/tfjs-backend-wasm/src/kernels/BatchToSpaceND.ts b/tfjs-backend-wasm/src/kernels/BatchToSpaceND.ts
@@ -0,0 +1,68 @@
+/**
+ * @license
+ * Copyright 2021 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 {backend_util, BatchToSpaceND, BatchToSpaceNDAttrs, BatchToSpaceNDInputs, KernelConfig, KernelFunc} from '@tensorflow/tfjs-core';
+
+import {BackendWasm} from '../backend_wasm';
+
+import {reshape} from './Reshape';
+import {slice} from './Slice';
+import {transpose} from './Transpose';
+
+function batchToSpaceND(args: {
+  inputs: BatchToSpaceNDInputs,
+  backend: BackendWasm,
+  attrs: BatchToSpaceNDAttrs
+}) {
+  const {inputs, backend, attrs} = args;
+  const {x} = inputs;
+  const {blockShape, crops} = attrs;
+
+  const prod = blockShape.reduce((a, b) => a * b);
+
+  const reshaped = backend_util.getReshaped(x.shape, blockShape, prod);
+  const permuted = backend_util.getPermuted(reshaped.length, blockShape.length);
+  const reshapedPermuted =
+      backend_util.getReshapedPermuted(x.shape, blockShape, prod);
+  const sliceBeginCoords =
+      backend_util.getSliceBeginCoords(crops, blockShape.length);
+  const sliceSize =
+      backend_util.getSliceSize(reshapedPermuted, crops, blockShape.length);
+
+  const xReshaped = reshape({inputs: {x}, backend, attrs: {shape: reshaped}});
+  const xTransposed =
+      transpose({inputs: {x: xReshaped}, backend, attrs: {perm: permuted}});
+  const xTransposedReshaped = reshape(
+      {inputs: {x: xTransposed}, backend, attrs: {shape: reshapedPermuted}});
+  const result = slice({
+    inputs: {x: xTransposedReshaped},
+    backend,
+    attrs: {begin: sliceBeginCoords, size: sliceSize}
+  });
+
+  backend.disposeData(xReshaped.dataId);
+  backend.disposeData(xTransposed.dataId);
+  backend.disposeData(xReshaped.dataId);
+
+  return result;
+}
+
+export const batchToSpaceNDConfig: KernelConfig = {
+  kernelName: BatchToSpaceND,
+  backendName: 'wasm',
+  kernelFunc: batchToSpaceND as {} as KernelFunc
+};
diff --git a/tfjs-backend-wasm/src/kernels/Reshape.ts b/tfjs-backend-wasm/src/kernels/Reshape.ts
@@ -15,18 +15,15 @@
  * =============================================================================
  */
 
-import {KernelConfig, NamedAttrMap, NamedTensorInfoMap, Reshape, ReshapeAttrs, ReshapeInputs, util} from '@tensorflow/tfjs-core';
+import {KernelConfig, KernelFunc, Reshape, ReshapeAttrs, ReshapeInputs, util} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
-export function reshape(args: {
-  inputs: NamedTensorInfoMap,
-  attrs: NamedAttrMap,
-  backend: BackendWasm
-}) {
+export function reshape(
+    args: {inputs: ReshapeInputs, attrs: ReshapeAttrs, backend: BackendWasm}) {
   const {inputs, attrs} = args;
-  const {x} = inputs as {} as ReshapeInputs;
-  const {shape} = attrs as {} as ReshapeAttrs;
+  const {x} = inputs;
+  const {shape} = attrs;
 
   const xSize = util.sizeFromShape(x.shape);
   const $shape = util.inferFromImplicitShape(shape, xSize);
@@ -44,5 +41,5 @@ export function reshape(args: {
 export const reshapeConfig: KernelConfig = {
   kernelName: Reshape,
   backendName: 'wasm',
-  kernelFunc: reshape,
+  kernelFunc: reshape as {} as KernelFunc
 };
diff --git a/tfjs-backend-wasm/src/kernels/SpaceToBatchND.ts b/tfjs-backend-wasm/src/kernels/SpaceToBatchND.ts
@@ -0,0 +1,86 @@
+/**
+ * @license
+ * Copyright 2021 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 {backend_util, KernelConfig, KernelFunc, ReshapeAttrs, ReshapeInputs, SpaceToBatchND, SpaceToBatchNDAttrs, SpaceToBatchNDInputs, TensorInfo, TransposeAttrs, TransposeInputs, util} from '@tensorflow/tfjs-core';
+
+import {BackendWasm} from '../backend_wasm';
+
+import {padV2Config} from './PadV2';
+import {reshape} from './Reshape';
+import {transpose} from './Transpose';
+
+function spaceToBatchND(args: {
+  inputs: SpaceToBatchNDInputs,
+  backend: BackendWasm,
+  attrs: SpaceToBatchNDAttrs
+}) {
+  const {inputs, backend, attrs} = args;
+  const {x} = inputs;
+  const {blockShape, paddings} = attrs;
+
+  const prod = util.sizeFromShape(blockShape);
+
+  const completePaddings: Array<[number, number]> = [[0, 0]];
+  completePaddings.push(...(paddings as Array<[number, number]>));
+
+  for (let i = 1 + blockShape.length; i < x.shape.length; ++i) {
+    completePaddings.push([0, 0]);
+  }
+
+  const paddedX = padV2Config.kernelFunc({
+    inputs: {x},
+    backend,
+    attrs: {paddings: completePaddings, constantValue: 0}
+  }) as TensorInfo;
+
+  const reshapedPaddedShape =
+      backend_util.getReshaped(paddedX.shape, blockShape, prod, false);
+
+  const permutedReshapedPaddedPermutation = backend_util.getPermuted(
+      reshapedPaddedShape.length, blockShape.length, false);
+
+  const flattenShape =
+      backend_util.getReshapedPermuted(paddedX.shape, blockShape, prod, false);
+
+  const reshapeInputs: ReshapeInputs = {x: paddedX};
+  const reshapeAttrs: ReshapeAttrs = {shape: reshapedPaddedShape};
+  const paddedXReshaped =
+      reshape({inputs: reshapeInputs, backend, attrs: reshapeAttrs});
+
+  const transposeInputs: TransposeInputs = {x: paddedXReshaped};
+  const transposeAttrs:
+      TransposeAttrs = {perm: permutedReshapedPaddedPermutation};
+  const paddedXT =
+      transpose({inputs: transposeInputs, backend, attrs: transposeAttrs});
+
+  const resultReshapeInputs: ReshapeInputs = {x: paddedXT};
+  const resultReshapeAttrs: ReshapeAttrs = {shape: flattenShape};
+  const result = reshape(
+      {inputs: resultReshapeInputs, backend, attrs: resultReshapeAttrs});
+
+  backend.disposeData(paddedX.dataId);
+  backend.disposeData(paddedXReshaped.dataId);
+  backend.disposeData(paddedXT.dataId);
+
+  return result;
+}
+
+export const spaceToBatchNDConfig: KernelConfig = {
+  kernelName: SpaceToBatchND,
+  backendName: 'wasm',
+  kernelFunc: spaceToBatchND as {} as KernelFunc
+};
diff --git a/tfjs-backend-wasm/src/register_all_kernels.ts b/tfjs-backend-wasm/src/register_all_kernels.ts
@@ -28,6 +28,7 @@ import {anyConfig} from './kernels/Any';
 import {argMaxConfig} from './kernels/ArgMax';
 import {avgPoolConfig} from './kernels/AvgPool';
 import {batchMatMulConfig} from './kernels/BatchMatMul';
+import {batchToSpaceNDConfig} from './kernels/BatchToSpaceND';
 import {castConfig} from './kernels/Cast';
 import {ceilConfig} from './kernels/Ceil';
 import {clipByValueConfig} from './kernels/ClipByValue';
@@ -96,6 +97,7 @@ import {sigmoidConfig} from './kernels/Sigmoid';
 import {sinConfig} from './kernels/Sin';
 import {sliceConfig} from './kernels/Slice';
 import {softmaxConfig} from './kernels/Softmax';
+import {spaceToBatchNDConfig} from './kernels/SpaceToBatchND';
 import {splitVConfig} from './kernels/SplitV';
 import {sqrtConfig} from './kernels/Sqrt';
 import {squareConfig} from './kernels/Square';
@@ -123,6 +125,7 @@ const kernelConfigs: KernelConfig[] = [
   argMaxConfig,
   avgPoolConfig,
   batchMatMulConfig,
+  batchToSpaceNDConfig,
   castConfig,
   ceilConfig,
   clipByValueConfig,
@@ -192,6 +195,7 @@ const kernelConfigs: KernelConfig[] = [
   sinConfig,
   sliceConfig,
   softmaxConfig,
+  spaceToBatchNDConfig,
   splitVConfig,
   sqrtConfig,
   squareConfig,

diff --git a/tfjs-backend-wasm/src/setup_test.ts b/tfjs-backend-wasm/src/setup_test.ts
@@ -424,7 +424,9 @@ const TEST_FILTERS: TestFilter[] = [
       'ignores NaNs'  // Doesn't yet ignore NaN
     ]
   },
-  {include: 'image.transform'}
+  {include: 'image.transform'},
+  {include: 'batchToSpaceND'},
+  {include: 'spaceToBatchND'},
 ];
 
 const customInclude = (testName: string) => {