Modularize spaceToBatchND.

tfjs-core/src/gradients/BatchToSpaceND_grad.ts

@@ -17,7 +17,7 @@
 
 import {BatchToSpaceND, BatchToSpaceNDAttrs} from '../kernel_names';
 import {GradConfig, NamedAttrMap} from '../kernel_registry';
-import {spaceToBatchND} from '../ops/array_ops';
+import {spaceToBatchND} from '../ops/space_to_batch_nd';
 import {Tensor} from '../tensor';
 
 export const batchToSpaceNDGradConfig: GradConfig = {

tfjs-core/src/gradients/SpaceToBatchND_grad.ts

@@ -0,0 +1,29 @@
+/**
+ * @license
+ * Copyright 2020 Google Inc. 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 {SpaceToBatchND, SpaceToBatchNDAttrs} from '../kernel_names';
+import {GradConfig, NamedAttrMap} from '../kernel_registry';
+import {batchToSpaceND} from '../ops/batch_to_space_nd';
+import {Tensor} from '../tensor';
+
+export const spaceToBatchNDGradConfig: GradConfig = {
+  kernelName: SpaceToBatchND,
+  gradFunc: (dy: Tensor, saved: Tensor[], attrs: NamedAttrMap) => {
+    const {blockShape, paddings} = attrs as {} as SpaceToBatchNDAttrs;
+    return {x: () => batchToSpaceND(dy, blockShape, paddings)};
+  }
+};

tfjs-core/src/kernel_names.ts

@@ -221,6 +221,13 @@ export interface PadV2Attrs {
   constantValue: number;
 }
 
+export const SpaceToBatchND = 'SpaceToBatchND';
+export type SpaceToBatchNDInputs = Pick<NamedTensorInfoMap, 'x'>;
+export interface SpaceToBatchNDAttrs {
+  blockShape: number[];
+  paddings: number[][];
+}
+
 export const SplitV = 'SplitV';
 export type SplitVInputs = Pick<NamedTensorInfoMap, 'x'>;
 export interface SplitVAttrs {

tfjs-core/src/ops/array_ops.ts

@@ -159,91 +159,6 @@ function stack_<T extends Tensor>(
   return concat(expandedTensors, axis);
 }
 
-/**
- * This operation divides "spatial" dimensions `[1, ..., M]` of the input into
- * a grid of blocks of shape `blockShape`, and interleaves these blocks with
- * the "batch" dimension (0) such that in the output, the spatial
- * dimensions `[1, ..., M]` correspond to the position within the grid,
- * and the batch dimension combines both the position within a spatial block
- * and the original batch position. Prior to division into blocks,
- * the spatial dimensions of the input are optionally zero padded
- * according to `paddings`. See below for a precise description.
- *
- * ```js
- * const x = tf.tensor4d([1, 2, 3, 4], [1, 2, 2, 1]);
- * const blockShape = [2, 2];
- * const paddings = [[0, 0], [0, 0]];
- *
- * x.spaceToBatchND(blockShape, paddings).print();
- * ```
- *
- * @param x A `tf.Tensor`. N-D with `x.shape` = `[batch] + spatialShape +
- * remainingShape`, where spatialShape has `M` dimensions.
- * @param blockShape A 1-D array. Must have shape `[M]`, all values must
- * be >= 1.
- * @param paddings A 2-D array. Must have shape `[M, 2]`, all values must be >=
- *     0. `paddings[i] = [padStart, padEnd]` specifies the amount to zero-pad
- * from input dimension `i + 1`, which corresponds to spatial dimension `i`. It
- * is required that
- * `(inputShape[i + 1] + padStart + padEnd) % blockShape[i] === 0`
- *
- * This operation is equivalent to the following steps:
- *
- * 1. Zero-pad the start and end of dimensions `[1, ..., M]` of the input
- * according to `paddings` to produce `padded` of shape paddedShape.
- *
- * 2. Reshape `padded` to `reshapedPadded` of shape:
- * `[batch] + [paddedShape[1] / blockShape[0], blockShape[0], ...,
- * paddedShape[M] / blockShape[M-1], blockShape[M-1]] + remainingShape`
- *
- * 3. Permute dimensions of `reshapedPadded` to produce `permutedReshapedPadded`
- * of shape: `blockShape + [batch] + [paddedShape[1] / blockShape[0], ...,
- * paddedShape[M] / blockShape[M-1]] + remainingShape`
- *
- * 4. Reshape `permutedReshapedPadded` to flatten `blockShape` into the
- * batch dimension, producing an output tensor of shape:
- * `[batch * prod(blockShape)] + [paddedShape[1] / blockShape[0], ...,
- * paddedShape[M] / blockShape[M-1]] + remainingShape`
- */
-/** @doc {heading: 'Tensors', subheading: 'Transformations'} */
-function spaceToBatchND_<T extends Tensor>(
-    x: T|TensorLike, blockShape: number[], paddings: number[][]): T {
-  const $x = convertToTensor(x, 'x', 'spaceToBatchND');
-
-  util.assert(
-      $x.rank >= 1 + blockShape.length,
-      () => `input rank ${$x.rank} should be > than [blockShape] ${
-          blockShape.length}`);
-
-  util.assert(
-      paddings.length === blockShape.length,
-      () => `paddings.shape[0] ${
-          paddings.length} must be equal to [blockShape] ${blockShape.length}`);
-
-  util.assert(
-      $x.shape.reduce(
-          (a, b, i) => {
-            if (i > 0 && i <= blockShape.length) {
-              return a &&
-                  ((b + paddings[i - 1][0] + paddings[i - 1][1]) %
-                       blockShape[i - 1] ===
-                   0);
-            }
-            return a;
-          },
-          true),
-      () => `input spatial dimensions ${$x.shape.slice(1)} with paddings ${
-          paddings.toString()} must be divisible by blockShapes ${
-          blockShape.toString()}`);
-
-  const grad = (dy: T) => {
-    return {$x: () => dy.batchToSpaceND(blockShape, paddings) as T};
-  };
-
-  return ENGINE.runKernelFunc(
-      backend => backend.spaceToBatchND($x, blockShape, paddings), {$x}, grad);
-}
-
 /**
  * Unstacks a `tf.Tensor` of rank-`R` into a list of rank-`(R-1)` `tf.Tensor`s.
  *
@@ -548,7 +463,6 @@ export const cumsum = op({cumsum_});
 export const depthToSpace = op({depthToSpace_});
 export const expandDims = op({expandDims_});
 export const reshape = op({reshape_});
-export const spaceToBatchND = op({spaceToBatchND_});
 export const squeeze = op({squeeze_});
 export const stack = op({stack_});
 export const unstack = op({unstack_});

tfjs-core/src/ops/ops.ts

@@ -63,6 +63,7 @@ export {randomGamma} from './random_gamma';
 export {randomNormal} from './random_normal';
 export {randomUniform} from './random_uniform';
 export {separableConv2d} from './separable_conv2d';
+export {spaceToBatchND} from './space_to_batch_nd';
 export {split} from './split';
 export {square} from './square';
 export {squaredDifference} from './squared_difference';

tfjs-core/src/ops/pool.ts

@@ -22,10 +22,10 @@ import {convertToTensor} from '../tensor_util_env';
 import {TensorLike} from '../types';
 import * as util from '../util';
 
-import {spaceToBatchND} from './array_ops';
 import {batchToSpaceND} from './batch_to_space_nd';
 import * as conv_util from './conv_util';
 import {op} from './operation';
+import {spaceToBatchND} from './space_to_batch_nd';
 
 /**
  * Computes the 2D max pooling of an image.

tfjs-core/src/ops/space_to_batch_nd.ts

@@ -0,0 +1,117 @@
+/**
+ * @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 {ENGINE, ForwardFunc} from '../engine';
+import {SpaceToBatchND, SpaceToBatchNDAttrs, SpaceToBatchNDInputs} from '../kernel_names';
+import {NamedAttrMap} from '../kernel_registry';
+import {Tensor} from '../tensor';
+import {NamedTensorMap} from '../tensor_types';
+import {convertToTensor} from '../tensor_util_env';
+import {TensorLike} from '../types';
+import * as util from '../util';
+
+import {op} from './operation';
+
+/**
+ * This operation divides "spatial" dimensions `[1, ..., M]` of the input into
+ * a grid of blocks of shape `blockShape`, and interleaves these blocks with
+ * the "batch" dimension (0) such that in the output, the spatial
+ * dimensions `[1, ..., M]` correspond to the position within the grid,
+ * and the batch dimension combines both the position within a spatial block
+ * and the original batch position. Prior to division into blocks,
+ * the spatial dimensions of the input are optionally zero padded
+ * according to `paddings`. See below for a precise description.
+ *
+ * ```js
+ * const x = tf.tensor4d([1, 2, 3, 4], [1, 2, 2, 1]);
+ * const blockShape = [2, 2];
+ * const paddings = [[0, 0], [0, 0]];
+ *
+ * x.spaceToBatchND(blockShape, paddings).print();
+ * ```
+ *
+ * @param x A `tf.Tensor`. N-D with `x.shape` = `[batch] + spatialShape +
+ * remainingShape`, where spatialShape has `M` dimensions.
+ * @param blockShape A 1-D array. Must have shape `[M]`, all values must
+ * be >= 1.
+ * @param paddings A 2-D array. Must have shape `[M, 2]`, all values must be >=
+ *     0. `paddings[i] = [padStart, padEnd]` specifies the amount to zero-pad
+ * from input dimension `i + 1`, which corresponds to spatial dimension `i`. It
+ * is required that
+ * `(inputShape[i + 1] + padStart + padEnd) % blockShape[i] === 0`
+ *
+ * This operation is equivalent to the following steps:
+ *
+ * 1. Zero-pad the start and end of dimensions `[1, ..., M]` of the input
+ * according to `paddings` to produce `padded` of shape paddedShape.
+ *
+ * 2. Reshape `padded` to `reshapedPadded` of shape:
+ * `[batch] + [paddedShape[1] / blockShape[0], blockShape[0], ...,
+ * paddedShape[M] / blockShape[M-1], blockShape[M-1]] + remainingShape`
+ *
+ * 3. Permute dimensions of `reshapedPadded` to produce `permutedReshapedPadded`
+ * of shape: `blockShape + [batch] + [paddedShape[1] / blockShape[0], ...,
+ * paddedShape[M] / blockShape[M-1]] + remainingShape`
+ *
+ * 4. Reshape `permutedReshapedPadded` to flatten `blockShape` into the
+ * batch dimension, producing an output tensor of shape:
+ * `[batch * prod(blockShape)] + [paddedShape[1] / blockShape[0], ...,
+ * paddedShape[M] / blockShape[M-1]] + remainingShape`
+ */
+/** @doc {heading: 'Tensors', subheading: 'Transformations'} */
+function spaceToBatchND_<T extends Tensor>(
+    x: T|TensorLike, blockShape: number[], paddings: number[][]): T {
+  const $x = convertToTensor(x, 'x', 'spaceToBatchND');
+
+  util.assert(
+      $x.rank >= 1 + blockShape.length,
+      () => `input rank ${$x.rank} should be > than [blockShape] ${
+          blockShape.length}`);
+
+  util.assert(
+      paddings.length === blockShape.length,
+      () => `paddings.shape[0] ${
+          paddings.length} must be equal to [blockShape] ${blockShape.length}`);
+
+  util.assert(
+      $x.shape.reduce(
+          (a, b, i) => {
+            if (i > 0 && i <= blockShape.length) {
+              return a &&
+                  ((b + paddings[i - 1][0] + paddings[i - 1][1]) %
+                       blockShape[i - 1] ===
+                   0);
+            }
+            return a;
+          },
+          true),
+      () => `input spatial dimensions ${$x.shape.slice(1)} with paddings ${
+          paddings.toString()} must be divisible by blockShapes ${
+          blockShape.toString()}`);
+
+  const forward: ForwardFunc<T> = backend =>
+      backend.spaceToBatchND($x, blockShape, paddings);
+
+  const inputs: SpaceToBatchNDInputs = {x: $x};
+  const attrs: SpaceToBatchNDAttrs = {blockShape, paddings};
+
+  return ENGINE.runKernelFunc(
+      forward, inputs as {} as NamedTensorMap, null /* gradient */,
+      SpaceToBatchND, attrs as {} as NamedAttrMap);
+}
+
+export const spaceToBatchND = op({spaceToBatchND_});