Move reshape() and cast() to kernel backends. (#893)

* wip

* wip

* wip

* cleanup

* Add missing file

Author: nkreeger

src/kernels/backend.ts

@@ -1,7 +1,6 @@
-
 /**
  * @license
- * Copyright 2017 Google Inc. All Rights Reserved.
+ * Copyright 2018 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
@@ -19,7 +18,7 @@
 import {Conv2DInfo} from '../ops/conv_util';
 // tslint:disable-next-line:max-line-length
 import {DataId, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D} from '../tensor';
-import {DataType, TypedArray} from '../types';
+import {DataType, Rank, ShapeMap, TypedArray} from '../types';
 
 // Required information for all backends.
 export interface BackendTimingInfo { kernelMs: number; }
@@ -158,6 +157,10 @@ export interface KernelBackend extends TensorStorage, BackendTimer {
   avgPool(x: Tensor4D, convInfo: Conv2DInfo): Tensor4D;
   avgPoolBackprop(dy: Tensor4D, x: Tensor4D, convInfo: Conv2DInfo): Tensor4D;
 
+  reshape<T extends Tensor, R extends Rank>(x: T, shape: ShapeMap[R]):
+      Tensor<R>;
+  cast<T extends Tensor>(x: T, dtype: DataType): T;
+
   tile<T extends Tensor>(x: T, reps: number[]): T;
 
   pad<T extends Tensor>(

src/kernels/backend_cpu.ts

@@ -32,6 +32,7 @@ import {DataType, DataTypeMap, Rank, TypedArray} from '../types';
 import * as util from '../util';
 
 import {BackendTimingInfo, KernelBackend} from './backend';
+import * as backend_util from './backend_util';
 
 export class MathBackendCPU implements KernelBackend {
   private data = new WeakMap<DataId, DataTypeMap[DataType]>();
@@ -1316,6 +1317,15 @@ export class MathBackendCPU implements KernelBackend {
     return dx.toTensor();
   }
 
+  cast<T extends Tensor<types.Rank>>(x: T, dtype: DataType): T {
+    return backend_util.castTensor(x, dtype, this);
+  }
+
+  reshape<T extends Tensor<types.Rank>, R extends types.Rank>(
+      x: T, shape: types.ShapeMap[R]): Tensor<R> {
+    return backend_util.reshapeTensor(x, shape);
+  }
+
   minPool(x: Tensor4D, convInfo: Conv2DInfo): Tensor4D {
     return this.pool(x, convInfo, 'min');
   }

src/kernels/backend_util.ts

@@ -0,0 +1,42 @@
+/**
+ * @license
+ * Copyright 2018 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 {Rank, Tensor, util} from '..';
+import {ArrayOps} from '../ops/array_ops';
+import {DataType, ShapeMap} from '../types';
+import {KernelBackend} from './backend';
+
+export function castTensor<T extends Tensor<Rank>>(
+    x: T, dtype: DataType, backend: KernelBackend): T {
+  if (!util.hasEncodingLoss(x.dtype, dtype)) {
+    // We don't change the underlying data, since we cast to higher
+    // precision.
+    return Tensor.make(x.shape, {dataId: x.dataId}, dtype) as T;
+  }
+  if (dtype === 'int32') {
+    return backend.int(x);
+  } else if (dtype === 'bool') {
+    return backend.notEqual(x, ArrayOps.scalar(0, x.dtype)) as T;
+  } else {
+    throw new Error(`Error in Cast: unknown dtype argument (${dtype})`);
+  }
+}
+
+export function reshapeTensor<T extends Tensor<Rank>, R extends Rank>(
+    x: T, shape: ShapeMap[R]): Tensor<R> {
+  return Tensor.make(shape, {dataId: x.dataId}, x.dtype);
+}

src/kernels/backend_webgl.ts

@@ -25,7 +25,9 @@ import {DataId, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D} from '../tensor'
 import * as types from '../types';
 import {DataType, DataTypeMap, RecursiveArray, TypedArray} from '../types';
 import * as util from '../util';
+
 import {KernelBackend} from './backend';
+import * as backend_util from './backend_util';
 import {ArgMinMaxProgram} from './webgl/argminmax_gpu';
 import {AvgPool2DBackpropProgram} from './webgl/avg_pool_backprop_gpu';
 import {BatchNormProgram} from './webgl/batchnorm_gpu';
@@ -813,6 +815,15 @@ export class MathBackendWebGL implements KernelBackend {
     return this.compileAndRun(avgPoolBackpropProgram, [dy], output) as Tensor4D;
   }
 
+  cast<T extends Tensor<types.Rank>>(x: T, dtype: DataType): T {
+    return backend_util.castTensor(x, dtype, this);
+  }
+
+  reshape<T extends Tensor<types.Rank>, R extends types.Rank>(
+      x: T, shape: types.ShapeMap[R]): Tensor<R> {
+    return backend_util.reshapeTensor(x, shape);
+  }
+
   resizeBilinear(
       x: Tensor4D, newHeight: number, newWidth: number,
       alignCorners: boolean): Tensor4D {

src/ops/array_ops.ts

@@ -16,7 +16,7 @@
  */
 
 import {doc} from '../doc';
-import {ForwardFunc} from '../engine';
+// import {ForwardFunc} from '../engine';
 import {ENV} from '../environment';
 // tslint:disable-next-line:max-line-length
 import {Scalar, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, TensorBuffer} from '../tensor';
@@ -606,7 +606,7 @@ export class ArrayOps {
       return {x: () => dy.reshape(x.shape)};
     };
     return ENV.engine.runKernel(
-        backend => Tensor.make(shape, {dataId: x.dataId}, x.dtype), {x}, grad);
+        backend => backend.reshape(x, shape), {x}, grad);
   }
 
   /**
@@ -640,24 +640,11 @@ export class ArrayOps {
   @doc({heading: 'Tensors', subheading: 'Transformations'})
   @operation
   static cast<T extends Tensor>(x: T, dtype: DataType): T {
-    const forw: ForwardFunc<T> = backend => {
-      if (!util.hasEncodingLoss(x.dtype, dtype)) {
-        // We don't change the underlying data, since we cast to higher
-        // precision.
-        return Tensor.make(x.shape, {dataId: x.dataId}, dtype) as T;
-      }
-      if (dtype === 'int32') {
-        return backend.int(x);
-      } else if (dtype === 'bool') {
-        return backend.notEqual(x, ArrayOps.scalar(0, x.dtype)) as T;
-      } else {
-        throw new Error(`Error in Cast: unknown dtype argument (${dtype})`);
-      }
-    };
     const grad = (dy: T) => {
       return {x: () => dy.clone()};
     };
-    return ENV.engine.runKernel(forw, {x}, grad) as T;
+    return ENV.engine.runKernel(backend => backend.cast(x, dtype), {x}, grad) as
+        T;
   }
 
   /**