backend: Add ops onesLike and zerosLike (tensorflow#1585)

- Also add CPU and WebGL backend implementations of these
  based on fill()
- This will be used in tfjs-node to bind to the libtensorflow
  implementations of onesLike and zerosLike to save memory
  copying between CPU and GPU, which will lead to performance
  improvements.

PERF

Author: caisq

src/kernels/backend.ts

@@ -589,6 +589,14 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer {
     throw new Error('Not yet implemented.');
   }
 
+  onesLike<R extends Rank>(x: Tensor<R>): Tensor<R> {
+    throw new Error('Not yet implemented');
+  }
+
+  zerosLike<R extends Rank>(x: Tensor<R>): Tensor<R> {
+    throw new Error('Not yet implemented');
+  }
+
   /**
    * Sets the data mover for this backend. Backends should use the mover to
    * move data from other backends to this backend.

src/kernels/backend_cpu.ts

@@ -3338,6 +3338,20 @@ export class MathBackendCPU implements KernelBackend {
     return Tensor.make(shape, {values}, dtype);
   }
 
+  onesLike<R extends Rank>(x: Tensor<R>): Tensor<R> {
+    if (x.dtype === 'string') {
+      throw new Error('onesLike is not supported for string tensors');
+    } else {
+      return this.fill(x.shape, 1, x.dtype);
+    }
+  }
+
+  zerosLike<R extends Rank>(x: Tensor<R>): Tensor<R> {
+    const values =
+        getArrayFromDType(x.dtype, sizeFromShape(x.shape)) as TypedArray;
+    return Tensor.make(x.shape, {values}, x.dtype);
+  }
+
   private scatter<R extends Rank>(
       indices: Tensor, updates: Tensor, shape: ShapeMap[R], outputSize: number,
       sliceSize: number, numUpdates: number, sliceRank: number,

src/kernels/backend_webgl.ts

@@ -2149,6 +2149,20 @@ export class MathBackendWebGL implements KernelBackend {
     }
   }
 
+  onesLike<R extends Rank>(x: Tensor<R>): Tensor<R> {
+    if (x.dtype === 'string') {
+      throw new Error('onesLike is not supported under string dtype');
+    } else {
+      // TODO(cais, smilkov): Add WebGL shader for onesLike:
+      //   https://github.com/tensorflow/tfjs/issues/1293
+      return this.fill(x.shape, 1, x.dtype);
+    }
+  }
+
+  zerosLike<R extends Rank>(x: Tensor<R>): Tensor<R> {
+    return this.fill(x.shape, x.dtype === 'string' ? '' : 0, x.dtype);
+  }
+
   private makeOutputArray<T extends Tensor>(shape: number[], dtype: DataType):
       T {
     return Tensor.make(shape, {}, dtype, this) as T;

src/ops/tensor_ops.ts

@@ -447,7 +447,7 @@ function fill<R extends Rank>(
 /** @doc {heading: 'Tensors', subheading: 'Creation'} */
 function onesLike_<T extends Tensor>(x: T|TensorLike): T {
   const $x = convertToTensor(x, 'x', 'onesLike');
-  return ones($x.shape, $x.dtype) as T;
+  return ENV.engine.runKernel(backend => backend.onesLike($x), {$x}, null) as T;
 }
 
 /**
@@ -464,7 +464,8 @@ function onesLike_<T extends Tensor>(x: T|TensorLike): T {
 /** @doc {heading: 'Tensors', subheading: 'Creation'} */
 function zerosLike_<T extends Tensor>(x: T|TensorLike): T {
   const $x = convertToTensor(x, 'x', 'zerosLike');
-  return zeros($x.shape, $x.dtype) as T;
+  return ENV.engine.runKernel(backend => backend.zerosLike($x), {$x}, null) as
+      T;
 }
 
 /**