Modularize depthToSpace.

tfjs-core/src/kernel_names.ts

@@ -107,6 +107,13 @@ export interface Conv3DBackpropInputAttrs {
   pad: 'valid'|'same';
 }
 
+export const DepthToSpace = 'DepthToSpace';
+export type DepthToSpaceInputs = Pick<NamedTensorInfoMap, 'x'>;
+export interface DepthToSpaceAttrs {
+  blockSize: number;
+  dataFormat: 'NHWC'|'NCHW';
+}
+
 export const DepthwiseConv2dNative = 'DepthwiseConv2dNative';
 export type DepthwiseConv2dNativeInputs =
     Pick<NamedTensorInfoMap, 'x'|'filter'>;

tfjs-core/src/ops/array_ops.ts

@@ -16,9 +16,9 @@
  */
 
 import {ENGINE} from '../engine';
-import {Tensor, Tensor4D, TensorBuffer} from '../tensor';
+import {Tensor, TensorBuffer} from '../tensor';
 import {convertToTensor, convertToTensorArray} from '../tensor_util_env';
-import {DataType, DataTypeMap, Rank, ShapeMap, TensorLike, TensorLike4D} from '../types';
+import {DataType, DataTypeMap, Rank, ShapeMap, TensorLike} from '../types';
 import * as util from '../util';
 import {getAxesPermutation, getInnerMostAxes} from './axis_util';
 import {concat} from './concat';
@@ -271,73 +271,6 @@ function expandDims_<R2 extends Rank>(
   return reshape($x, newShape as ShapeMap[R2]);
 }
 
-/**
- * Rearranges data from depth into blocks of spatial data. More specifically,
- * this op outputs a copy of the input tensor where values from the `depth`
- * dimension are moved in spatial blocks to the `height` and `width` dimensions.
- * The attr `blockSize` indicates the input block size and how the data is
- * moved.
- *
- *  - Chunks of data of size `blockSize * blockSize` from depth are rearranged
- * into non-overlapping blocks of size `blockSize x blockSize`
- *
- *  - The width the output tensor is `inputWidth * blockSize`, whereas the
- * height is `inputHeight * blockSize`
- *
- *  - The Y, X coordinates within each block of the output image are determined
- * by the high order component of the input channel index
- *
- *  - The depth of the input tensor must be divisible by `blockSize *
- * blockSize`
- *
- * The `dataFormat` attr specifies the layout of the input and output tensors
- * with the following options: "NHWC": [ `batch, height, width, channels` ]
- * "NCHW": [ `batch, channels, height, width` ]
- *
- * ```js
- * const x = tf.tensor4d([1, 2, 3, 4], [1, 1, 1, 4]);
- * const blockSize = 2;
- * const dataFormat = "NHWC";
- *
- * tf.depthToSpace(x, blockSize, dataFormat).print();
- * ```
- *
- * @param x The input tensor of rank 4
- * @param blockSIze  An `int` that is `>= 2`. The size of the spatial block
- * @param dataFormat An optional string from: "NHWC", "NCHW". Defaults to "NHWC"
- */
-/** @doc {heading: 'Tensors', subheading: 'Transformations'} */
-function depthToSpace_(
-    x: Tensor4D|TensorLike4D, blockSize: number,
-    dataFormat: 'NHWC'|'NCHW' = 'NHWC'): Tensor4D {
-  const $x = convertToTensor(x, 'x', 'depthToSpace') as Tensor4D;
-
-  const inputHeight = (dataFormat === 'NHWC') ? $x.shape[1] : $x.shape[2];
-  const inputWidth = (dataFormat === 'NHWC') ? $x.shape[2] : $x.shape[3];
-  const inputDepth = (dataFormat === 'NHWC') ? $x.shape[3] : $x.shape[1];
-
-  util.assert(
-      inputHeight * blockSize >= 0,
-      () => `Negative dimension size caused by overflow when multiplying
-      ${inputHeight} and ${blockSize}  for depthToSpace with input shape
-      ${$x.shape}`);
-
-  util.assert(
-      inputWidth * blockSize >= 0,
-      () => `Negative dimension size caused by overflow when multiplying
-      ${inputWidth} and ${blockSize} for depthToSpace with input shape
-          ${$x.shape}`);
-
-  util.assert(
-      (inputDepth % (blockSize * blockSize) === 0),
-      () => `Dimension size must be evenly divisible by ${
-          blockSize * blockSize} but is ${
-          inputDepth} for depthToSpace with input shape ${$x.shape}`);
-
-  return ENGINE.runKernelFunc(
-      backend => backend.depthToSpace($x, blockSize, dataFormat), {$x});
-}
-
 /**
  * Computes the difference between two lists of numbers.
  *
@@ -460,7 +393,6 @@ export {
 
 export const cast = op({cast_});
 export const cumsum = op({cumsum_});
-export const depthToSpace = op({depthToSpace_});
 export const expandDims = op({expandDims_});
 export const reshape = op({reshape_});
 export const squeeze = op({squeeze_});

tfjs-core/src/ops/array_ops_test.ts

@@ -3100,67 +3100,6 @@ describeWithFlags('batchToSpaceND X spaceToBatchND', ALL_ENVS, () => {
   });
 });
 
-describeWithFlags('depthToSpace', ALL_ENVS, () => {
-  it('tensor4d, input shape=[1, 1, 1, 4], blockSize=2, format=NHWC',
-     async () => {
-       const t = tf.tensor4d([[[[1, 2, 3, 4]]]]);
-       const blockSize = 2;
-       const dataFormat = 'NHWC';
-
-       const res = tf.depthToSpace(t, blockSize, dataFormat);
-       expect(res.shape).toEqual([1, 2, 2, 1]);
-       expectArraysClose(await res.data(), [1, 2, 3, 4]);
-     });
-
-  it('tensor4d, input shape=[1, 1, 1, 12], blockSize=2, format=NHWC',
-     async () => {
-       const t = tf.tensor4d([[[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]]]]);
-       const blockSize = 2;
-       const dataFormat = 'NHWC';
-
-       const res = tf.depthToSpace(t, blockSize, dataFormat);
-       expect(res.shape).toEqual([1, 2, 2, 3]);
-       expectArraysClose(
-           await res.data(), [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]);
-     });
-
-  it('tensor4d, input shape=[1, 2, 2, 4], blockSize=2, format=NHWC',
-     async () => {
-       const t = tf.tensor4d([
-         [[[1, 2, 3, 4], [5, 6, 7, 8]], [[9, 10, 11, 12], [13, 14, 15, 16]]]
-       ]);
-       const blockSize = 2;
-       const dataFormat = 'NHWC';
-
-       const res = tf.depthToSpace(t, blockSize, dataFormat);
-       expect(res.shape).toEqual([1, 4, 4, 1]);
-       expectArraysClose(
-           await res.data(),
-           [1, 2, 5, 6, 3, 4, 7, 8, 9, 10, 13, 14, 11, 12, 15, 16]);
-     });
-
-  it('throws when depth not divisible by blockSize * blockSize', () => {
-    const t = tf.tensor4d([1, 2, 3, 4], [1, 1, 1, 4]);
-    const blockSize = 3;
-
-    expect(() => tf.depthToSpace(t, blockSize))
-        .toThrowError(`Dimension size must be evenly divisible by ${
-            blockSize * blockSize} but is ${
-            t.shape[3]} for depthToSpace with input shape ${t.shape}`);
-  });
-});
-
-describeWithFlags('depthToSpace', BROWSER_ENVS, () => {
-  it('throws when blocksize < 2', () => {
-    const t = tf.tensor4d([1, 2, 3, 4], [1, 1, 1, 4]);
-    const blockSize = 1;
-
-    expect(() => tf.depthToSpace(t, blockSize))
-        .toThrowError(
-            `blockSize should be > 1 for depthToSpace, but was: ${blockSize}`);
-  });
-});
-
 describeWithFlags('setdiff1dAsync', ALL_ENVS, () => {
   it('1d int32 tensor', async () => {
     const x = tf.tensor1d([1, 2, 3, 4], 'int32');

tfjs-core/src/ops/depth_to_space.ts

@@ -0,0 +1,103 @@
+/**
+ * @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 {ENGINE, ForwardFunc} from '../engine';
+import {DepthToSpace, DepthToSpaceAttrs, DepthToSpaceInputs} from '../kernel_names';
+import {NamedAttrMap} from '../kernel_registry';
+import {Tensor4D} from '../tensor';
+import {NamedTensorMap} from '../tensor_types';
+import {convertToTensor} from '../tensor_util_env';
+import {TensorLike4D} from '../types';
+import * as util from '../util';
+
+import {op} from './operation';
+
+/**
+ * Rearranges data from depth into blocks of spatial data. More specifically,
+ * this op outputs a copy of the input tensor where values from the `depth`
+ * dimension are moved in spatial blocks to the `height` and `width` dimensions.
+ * The attr `blockSize` indicates the input block size and how the data is
+ * moved.
+ *
+ *  - Chunks of data of size `blockSize * blockSize` from depth are rearranged
+ * into non-overlapping blocks of size `blockSize x blockSize`
+ *
+ *  - The width the output tensor is `inputWidth * blockSize`, whereas the
+ * height is `inputHeight * blockSize`
+ *
+ *  - The Y, X coordinates within each block of the output image are determined
+ * by the high order component of the input channel index
+ *
+ *  - The depth of the input tensor must be divisible by `blockSize *
+ * blockSize`
+ *
+ * The `dataFormat` attr specifies the layout of the input and output tensors
+ * with the following options: "NHWC": [ `batch, height, width, channels` ]
+ * "NCHW": [ `batch, channels, height, width` ]
+ *
+ * ```js
+ * const x = tf.tensor4d([1, 2, 3, 4], [1, 1, 1, 4]);
+ * const blockSize = 2;
+ * const dataFormat = "NHWC";
+ *
+ * tf.depthToSpace(x, blockSize, dataFormat).print();
+ * ```
+ *
+ * @param x The input tensor of rank 4
+ * @param blockSIze  An `int` that is `>= 2`. The size of the spatial block
+ * @param dataFormat An optional string from: "NHWC", "NCHW". Defaults to "NHWC"
+ */
+/** @doc {heading: 'Tensors', subheading: 'Transformations'} */
+function depthToSpace_(
+    x: Tensor4D|TensorLike4D, blockSize: number,
+    dataFormat: 'NHWC'|'NCHW' = 'NHWC'): Tensor4D {
+  const $x = convertToTensor(x, 'x', 'depthToSpace') as Tensor4D;
+
+  const inputHeight = (dataFormat === 'NHWC') ? $x.shape[1] : $x.shape[2];
+  const inputWidth = (dataFormat === 'NHWC') ? $x.shape[2] : $x.shape[3];
+  const inputDepth = (dataFormat === 'NHWC') ? $x.shape[3] : $x.shape[1];
+
+  util.assert(
+      inputHeight * blockSize >= 0,
+      () => `Negative dimension size caused by overflow when multiplying
+    ${inputHeight} and ${blockSize}  for depthToSpace with input shape
+    ${$x.shape}`);
+
+  util.assert(
+      inputWidth * blockSize >= 0,
+      () => `Negative dimension size caused by overflow when multiplying
+    ${inputWidth} and ${blockSize} for depthToSpace with input shape
+        ${$x.shape}`);
+
+  util.assert(
+      (inputDepth % (blockSize * blockSize) === 0),
+      () => `Dimension size must be evenly divisible by ${
+          blockSize * blockSize} but is ${
+          inputDepth} for depthToSpace with input shape ${$x.shape}`);
+
+  const forward: ForwardFunc<Tensor4D> = backend =>
+      backend.depthToSpace($x, blockSize, dataFormat);
+
+  const inputs: DepthToSpaceInputs = {x: $x};
+  const attrs: DepthToSpaceAttrs = {blockSize, dataFormat};
+
+  return ENGINE.runKernelFunc(
+      forward, inputs as {} as NamedTensorMap, null /* gradient */,
+      DepthToSpace, attrs as {} as NamedAttrMap);
+}
+
+export const depthToSpace = op({depthToSpace_});

tfjs-core/src/ops/depth_to_space_test.ts

@@ -0,0 +1,81 @@
+/**
+ * @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 * as tf from '../index';
+import {ALL_ENVS, BROWSER_ENVS, describeWithFlags} from '../jasmine_util';
+import {expectArraysClose} from '../test_util';
+
+describeWithFlags('depthToSpace', ALL_ENVS, () => {
+  it('tensor4d, input shape=[1, 1, 1, 4], blockSize=2, format=NHWC',
+     async () => {
+       const t = tf.tensor4d([[[[1, 2, 3, 4]]]]);
+       const blockSize = 2;
+       const dataFormat = 'NHWC';
+
+       const res = tf.depthToSpace(t, blockSize, dataFormat);
+       expect(res.shape).toEqual([1, 2, 2, 1]);
+       expectArraysClose(await res.data(), [1, 2, 3, 4]);
+     });
+
+  it('tensor4d, input shape=[1, 1, 1, 12], blockSize=2, format=NHWC',
+     async () => {
+       const t = tf.tensor4d([[[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]]]]);
+       const blockSize = 2;
+       const dataFormat = 'NHWC';
+
+       const res = tf.depthToSpace(t, blockSize, dataFormat);
+       expect(res.shape).toEqual([1, 2, 2, 3]);
+       expectArraysClose(
+           await res.data(), [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]);
+     });
+
+  it('tensor4d, input shape=[1, 2, 2, 4], blockSize=2, format=NHWC',
+     async () => {
+       const t = tf.tensor4d([
+         [[[1, 2, 3, 4], [5, 6, 7, 8]], [[9, 10, 11, 12], [13, 14, 15, 16]]]
+       ]);
+       const blockSize = 2;
+       const dataFormat = 'NHWC';
+
+       const res = tf.depthToSpace(t, blockSize, dataFormat);
+       expect(res.shape).toEqual([1, 4, 4, 1]);
+       expectArraysClose(
+           await res.data(),
+           [1, 2, 5, 6, 3, 4, 7, 8, 9, 10, 13, 14, 11, 12, 15, 16]);
+     });
+
+  it('throws when depth not divisible by blockSize * blockSize', () => {
+    const t = tf.tensor4d([1, 2, 3, 4], [1, 1, 1, 4]);
+    const blockSize = 3;
+
+    expect(() => tf.depthToSpace(t, blockSize))
+        .toThrowError(`Dimension size must be evenly divisible by ${
+            blockSize * blockSize} but is ${
+            t.shape[3]} for depthToSpace with input shape ${t.shape}`);
+  });
+});
+
+describeWithFlags('depthToSpace', BROWSER_ENVS, () => {
+  it('throws when blocksize < 2', () => {
+    const t = tf.tensor4d([1, 2, 3, 4], [1, 1, 1, 4]);
+    const blockSize = 1;
+
+    expect(() => tf.depthToSpace(t, blockSize))
+        .toThrowError(
+            `blockSize should be > 1 for depthToSpace, but was: ${blockSize}`);
+  });
+});

tfjs-core/src/ops/ops.ts

@@ -35,6 +35,7 @@ export {conv2d} from './conv2d';
 export {conv2dTranspose} from './conv2d_transpose';
 export {conv3d} from './conv3d';
 export {conv3dTranspose} from './conv3d_transpose';
+export {depthToSpace} from './depth_to_space';
 export {depthwiseConv2d} from './depthwise_conv2d';
 export {diag} from './diag';
 export {div} from './div';