[core] Remove reshape from tensor / opHandler.

tfjs-backend-cpu/src/backend_cpu.ts

@@ -18,7 +18,7 @@
 import * as tf from '@tensorflow/tfjs-core';
 import {engine, env} from '@tensorflow/tfjs-core';
 import {backend_util, buffer, slice_util, util} from '@tensorflow/tfjs-core';
-import {BackendTimingInfo, DataStorage, DataType, DataValues, KernelBackend, max, NumericDataType, Rank, Scalar, ShapeMap, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, TensorBuffer, TypedArray, upcastType} from '@tensorflow/tfjs-core';
+import {BackendTimingInfo, DataStorage, DataType, DataValues, KernelBackend, max, NumericDataType, Rank, reshape, Scalar, ShapeMap, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, TensorBuffer, TypedArray, upcastType} from '@tensorflow/tfjs-core';
 import {kernel_impls} from '@tensorflow/tfjs-core';
 
 const nonMaxSuppressionV3Impl = kernel_impls.nonMaxSuppressionV3Impl;
@@ -2145,10 +2145,10 @@ export class MathBackendCPU extends KernelBackend {
     const flattenShape = backend_util.getReshapedPermuted(
         paddedX.shape, blockShape, prod, false);
 
-    return tf.transpose(
-                 paddedX.reshape(reshapedPaddedShape),
-                 permutedReshapedPaddedPermutation)
-               .reshape(flattenShape) as T;
+    const paddedXT = tf.transpose(
+        paddedX.reshape(reshapedPaddedShape),
+        permutedReshapedPaddedPermutation);
+    return reshape(paddedXT, flattenShape) as T;
   }
 
   maxPool(x: Tensor4D, convInfo: backend_util.Conv2DInfo): Tensor4D {

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

@@ -25,9 +25,9 @@ export function reshape(args: {
   backend: BackendWasm
 }) {
   const {inputs, attrs} = args;
-  const {tensor} = inputs as {} as ReshapeInputs;
+  const {x} = inputs as {} as ReshapeInputs;
   const {shape} = attrs as {} as ReshapeAttrs;
-  return {dataId: tensor.dataId, shape, dtype: tensor.dtype};
+  return {dataId: x.dataId, shape, dtype: x.dtype};
 }
 
 registerKernel({

tfjs-backend-webgl/src/backend_webgl.ts

@@ -19,7 +19,7 @@
 import './flags_webgl';
 
 import * as tf from '@tensorflow/tfjs-core';
-import {complex, DataId, div, engine, env, imag, max, MemoryInfo, range, real, RecursiveArray, scalar, softmax, tensor, tidy, TimingInfo, transpose} from '@tensorflow/tfjs-core';
+import {complex, DataId, div, engine, env, imag, max, MemoryInfo, range, real, RecursiveArray, reshape, scalar, softmax, tensor, tidy, TimingInfo, transpose} from '@tensorflow/tfjs-core';
 import {backend_util, buffer, kernel_impls, slice_util, util} from '@tensorflow/tfjs-core';
 import {DataStorage, DataType, KernelBackend, NumericDataType, Rank, Scalar, ShapeMap, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, TensorInfo, TypedArray, upcastType} from '@tensorflow/tfjs-core';
 
@@ -1020,10 +1020,10 @@ export class MathBackendWebGL extends KernelBackend {
     const flattenShape = backend_util.getReshapedPermuted(
         paddedX.shape, blockShape, prod, false);
 
-    return transpose(
-               paddedX.reshape(reshapedPaddedShape),
-               permutedReshapedPaddedPermutation)
-               .reshape(flattenShape) as T;
+    const paddedXT = transpose(
+        paddedX.reshape(reshapedPaddedShape),
+        permutedReshapedPaddedPermutation);
+    return reshape(paddedXT, flattenShape) as T;
   }
 
   private reduce(

tfjs-core/src/gradients/Reshape_grad.ts

@@ -21,9 +21,9 @@ import {Tensor} from '../tensor';
 
 export const reshapeGradConfig: GradConfig = {
   kernelName: Reshape,
-  inputsToSave: ['tensor'],
+  inputsToSave: ['x'],
   gradFunc: (dy: Tensor, saved: Tensor[]) => {
-    const [tensor] = saved;
-    return {tensor: () => reshape(dy, tensor.shape)};
+    const [x] = saved;
+    return {x: () => reshape(dy, x.shape)};
   }
 };

tfjs-core/src/kernel_names.ts

@@ -438,7 +438,7 @@ export const Relu = 'Relu';
 export type ReluInputs = Pick<NamedTensorInfoMap, 'x'>;
 
 export const Reshape = 'Reshape';
-export type ReshapeInputs = Pick<NamedTensorInfoMap, 'tensor'>;
+export type ReshapeInputs = Pick<NamedTensorInfoMap, 'x'>;
 export interface ReshapeAttrs {
   shape: number[];
 }

tfjs-core/src/ops/all.ts

@@ -70,7 +70,7 @@ function all_<T extends Tensor>(
     const res = backend.all($x, axes);
     if (keepDims) {
       const newShape = expandShapeToKeepDim(res.shape, origAxes);
-      return res.reshape(newShape) as T;
+      return res.reshape(newShape);
     }
     return res as T;
   };

tfjs-core/src/ops/any.ts

@@ -70,7 +70,7 @@ function any_<T extends Tensor>(
     const res = backend.any($x, axes);
     if (keepDims) {
       const newShape = expandShapeToKeepDim(res.shape, origAxes);
-      return res.reshape(newShape) as T;
+      return res.reshape(newShape);
     }
     return res as T;
   };

tfjs-core/src/ops/fused_ops.ts

@@ -256,7 +256,7 @@ function fusedMatMul_<T extends Tensor>({
       },
       inputs, grad, '_FusedMatMul', {transposeA, transposeB, activation},
       inputsToSave, outputsToSave);
-  return res.reshape(outShape) as T;
+  return res.reshape(outShape);
 }
 
 /**

tfjs-core/src/ops/reduction_ops.ts

@@ -221,7 +221,7 @@ function min_<T extends Tensor>(
   }, {x: $x}, grad, 'Min', {axes}, inputsToSave, outputsToSave);
   if (keepDims) {
     const newShape = axis_util.expandShapeToKeepDim(res.shape, origAxes);
-    res = res.reshape(newShape) as T;
+    res = res.reshape(newShape);
   }
   return res;
 }

tfjs-core/src/ops/reduction_ops_util.ts

@@ -24,10 +24,10 @@ import * as axis_util from './axis_util';
 export function gradForMinAndMax<T extends Tensor>(
     dy: T, y: T, xOrig: Tensor, origAxes: number[], permutedAxes: number[]) {
   if (y.rank < xOrig.rank) {
-    y = y.reshape(axis_util.expandShapeToKeepDim(y.shape, origAxes)) as T;
+    y = y.reshape(axis_util.expandShapeToKeepDim(y.shape, origAxes));
   }
   if (dy.rank < xOrig.rank) {
-    dy = dy.reshape(axis_util.expandShapeToKeepDim(dy.shape, origAxes)) as T;
+    dy = dy.reshape(axis_util.expandShapeToKeepDim(dy.shape, origAxes));
   }
   return {
     x: () => {

tfjs-core/src/ops/reshape.ts

@@ -52,17 +52,17 @@ import {op} from './operation';
  * @param shape An array of integers defining the output tensor shape.
  */
 /** @doc {heading: 'Tensors', subheading: 'Transformations'} */
-function reshape_<R2 extends Rank>(
-    x: Tensor|TensorLike, shape: ShapeMap[R2]): Tensor<R2> {
+function reshape_<R extends Rank>(
+    x: Tensor|TensorLike, shape: ShapeMap[R]): Tensor<R> {
   const $x = convertToTensor(x, 'x', 'reshape', null);
-  shape = util.inferFromImplicitShape(shape, $x.size) as ShapeMap[R2];
+  shape = util.inferFromImplicitShape(shape, $x.size) as ShapeMap[R];
   util.assert(
       $x.size === util.sizeFromShape(shape),
       () => 'new shape and old shape must have the same number of elements.');
 
-  const inputs: ReshapeInputs = {tensor: $x};
+  const inputs: ReshapeInputs = {x: $x};
   const attrs: ReshapeAttrs = {shape};
-  const forward: ForwardFunc<Tensor<R2>> =
+  const forward: ForwardFunc<Tensor<R>> =
       (backend: KernelBackend, save: GradSaveFunc) => {
         save([$x]);
         return backend.reshape($x, shape);

tfjs-core/src/ops/segment_ops.ts

@@ -137,7 +137,7 @@ function gather_<T extends Tensor>(
                 return res;
               },
               {x: $x, indices: $indices}, grad, 'Gather', {axis}))
-             .reshape(shapeInfo.outputShape) as T;
+      .reshape(shapeInfo.outputShape);
 }
 
 function arrayRange(start: number, stop: number): number[] {

tfjs-core/src/public/chained_ops/register_all_chained_ops.ts

@@ -64,6 +64,7 @@ import './pow';
 import './prelu';
 import './prod';
 import './relu';
+import './reshape';
 import './resize_bilinear';
 import './resize_nearest_neighbor';
 import './relu6';

tfjs-core/src/public/chained_ops/register_all_chained_ops_test.ts

@@ -74,6 +74,7 @@ const CHAINED_OPS = [
   'prelu',
   'prod',
   'relu',
+  'reshape',
   'resizeBilinear',
   'resizeNearestNeighbor',
   'relu6',

tfjs-core/src/public/chained_ops/reshape.ts

@@ -0,0 +1,30 @@
+/**
+ * @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 {reshape} from '../../ops/reshape';
+import {Tensor} from '../../tensor';
+import {Rank} from '../../types';
+
+declare module '../../tensor' {
+  interface Tensor<R extends Rank = Rank> {
+    reshape<T extends Tensor>(shape: number[]): T;
+  }
+}
+
+Tensor.prototype.reshape = function<T extends Tensor>(shape: number[]): T {
+  this.throwIfDisposed();
+  return reshape(this, shape) as T;
+};

tfjs-core/src/tensor.ts

@@ -174,7 +174,6 @@ export interface OpHandler {
       shape: ShapeMap[R], dtype: D,
       values?: DataTypeMap[D]): TensorBuffer<R, D>;
   print<T extends Tensor>(x: T, verbose: boolean): void;
-  reshape<R2 extends Rank>(x: Tensor, shape: ShapeMap[R2]): Tensor<R2>;
   clone<T extends Tensor>(x: T): T;
   gather<T extends Tensor>(x: T, indices: Tensor|TensorLike, axis: number): T;
   norm(
@@ -364,14 +363,14 @@ export class Tensor<R extends Rank = Rank> {
   asScalar(): Scalar {
     this.throwIfDisposed();
     util.assert(this.size === 1, () => 'The array must have only 1 element.');
-    return this.reshape<Rank.R0>([]);
+    return this.reshape([]);
   }
 
   /** Converts a `tf.Tensor` to a `tf.Tensor1D`. */
   /** @doc {heading: 'Tensors', subheading: 'Classes'} */
   as1D(): Tensor1D {
     this.throwIfDisposed();
-    return this.reshape<Rank.R1>([this.size]);
+    return this.reshape([this.size]);
   }
 
   /**
@@ -383,7 +382,7 @@ export class Tensor<R extends Rank = Rank> {
   /** @doc {heading: 'Tensors', subheading: 'Classes'} */
   as2D(rows: number, columns: number): Tensor2D {
     this.throwIfDisposed();
-    return this.reshape<Rank.R2>([rows, columns]);
+    return this.reshape([rows, columns]);
   }
 
   /**
@@ -396,7 +395,7 @@ export class Tensor<R extends Rank = Rank> {
   /** @doc {heading: 'Tensors', subheading: 'Classes'} */
   as3D(rows: number, columns: number, depth: number): Tensor3D {
     this.throwIfDisposed();
-    return this.reshape<Rank.R3>([rows, columns, depth]);
+    return this.reshape([rows, columns, depth]);
   }
 
   /**
@@ -410,7 +409,7 @@ export class Tensor<R extends Rank = Rank> {
   /** @doc {heading: 'Tensors', subheading: 'Classes'} */
   as4D(rows: number, columns: number, depth: number, depth2: number): Tensor4D {
     this.throwIfDisposed();
-    return this.reshape<Rank.R4>([rows, columns, depth, depth2]);
+    return this.reshape([rows, columns, depth, depth2]);
   }
 
   /**
@@ -427,7 +426,7 @@ export class Tensor<R extends Rank = Rank> {
       rows: number, columns: number, depth: number, depth2: number,
       depth3: number): Tensor5D {
     this.throwIfDisposed();
-    return this.reshape<Rank.R5>([rows, columns, depth, depth2, depth3]);
+    return this.reshape([rows, columns, depth, depth2, depth3]);
   }
 
   /**
@@ -584,18 +583,6 @@ export class Tensor<R extends Rank = Rank> {
     return opHandler.print(this, verbose);
   }
 
-  /**
-   * Reshapes the tensor into the provided shape.
-   * See `tf.reshape` for more details.
-   *
-   * @param newShape An array of integers defining the output tensor shape.
-   */
-  /** @doc {heading: 'Tensors', subheading: 'Classes'} */
-  reshape<R2 extends Rank>(newShape: ShapeMap[R2]): Tensor<R2> {
-    this.throwIfDisposed();
-    return opHandler.reshape(this, newShape);
-  }
-
   /**
    * Reshapes the tensor into the shape of the provided tensor.
    *
@@ -604,7 +591,7 @@ export class Tensor<R extends Rank = Rank> {
   /** @doc {heading: 'Tensors', subheading: 'Classes'} */
   reshapeAs<T extends Tensor>(x: T): T {
     this.throwIfDisposed();
-    return this.reshape(x.shape) as T;
+    return this.reshape(x.shape);
   }
 
   /** Returns a copy of the tensor. See `tf.clone` for details. */