Support complex type in onesLike. (tensorflow#1609)

FEATURE

Author: Lewuathe

src/ops/array_ops_test.ts

@@ -511,6 +511,16 @@ describeWithFlags('onesLike', ALL_ENVS, () => {
     expectArraysEqual(b, [1, 1, 1]);
   });
 
+  it('1D complex dtype', () => {
+    const real = tf.tensor1d([1, 2, 3], 'float32');
+    const imag = tf.tensor1d([1, 2, 3], 'float32');
+    const a = tf.complex(real, imag);
+    const b = tf.onesLike(a);
+    expect(b.dtype).toBe('complex64');
+    expect(b.shape).toEqual([3]);
+    expectArraysEqual(b, [1, 0, 1, 0, 1, 0]);
+  });
+
   it('2D default dtype', () => {
     const a = tf.tensor2d([1, 2, 3, 4], [2, 2]);
     const b = tf.onesLike(a);
@@ -543,6 +553,16 @@ describeWithFlags('onesLike', ALL_ENVS, () => {
     expectArraysEqual(b, [1, 1, 1, 1]);
   });
 
+  it('2D complex dtype', () => {
+    const real = tf.tensor2d([1, 2, 3, 4], [2, 2], 'float32');
+    const imag = tf.tensor2d([1, 2, 3, 4], [2, 2], 'float32');
+    const a = tf.complex(real, imag);
+    const b = tf.onesLike(a);
+    expect(b.dtype).toBe('complex64');
+    expect(b.shape).toEqual([2, 2]);
+    expectArraysEqual(b, [1, 0, 1, 0, 1, 0, 1, 0]);
+  });
+
   it('3D default dtype', () => {
     const a = tf.tensor3d([1, 2, 3, 4], [2, 2, 1]);
     const b = tf.onesLike(a);
@@ -575,6 +595,16 @@ describeWithFlags('onesLike', ALL_ENVS, () => {
     expectArraysEqual(b, [1, 1, 1, 1]);
   });
 
+  it('3D complex dtype', () => {
+    const real = tf.tensor3d([1, 2, 3, 4], [2, 2, 1], 'float32');
+    const imag = tf.tensor3d([1, 2, 3, 4], [2, 2, 1], 'float32');
+    const a = tf.complex(real, imag);
+    const b = tf.onesLike(a);
+    expect(b.dtype).toBe('complex64');
+    expect(b.shape).toEqual([2, 2, 1]);
+    expectArraysEqual(b, [1, 0, 1, 0, 1, 0, 1, 0]);
+  });
+
   it('4D default dtype', () => {
     const a = tf.tensor4d([1, 2, 3, 4], [2, 2, 1, 1]);
     const b = tf.onesLike(a);
@@ -615,6 +645,16 @@ describeWithFlags('onesLike', ALL_ENVS, () => {
     expectArraysClose(b, [1, 1, 1, 1]);
   });
 
+  it('4D complex dtype', () => {
+    const real = tf.tensor4d([1, 2, 3, 4], [2, 2, 1, 1], 'float32');
+    const imag = tf.tensor4d([1, 2, 3, 4], [2, 2, 1, 1], 'float32');
+    const a = tf.complex(real, imag);
+    const b = tf.onesLike(a);
+    expect(b.dtype).toBe('complex64');
+    expect(b.shape).toEqual([2, 2, 1, 1]);
+    expectArraysEqual(b, [1, 0, 1, 0, 1, 0, 1, 0]);
+  });
+
   it('5D float32 dtype', () => {
     const a = tf.tensor5d([1, 2, 3, 4], [1, 2, 2, 1, 1], 'float32');
     const b = tf.onesLike(a);
@@ -647,6 +687,16 @@ describeWithFlags('onesLike', ALL_ENVS, () => {
     expectArraysClose(b, [1, 1, 1, 1]);
   });
 
+  it('5D complex dtype', () => {
+    const real = tf.tensor5d([1, 2, 3, 4], [1, 2, 2, 1, 1], 'float32');
+    const imag = tf.tensor5d([1, 2, 3, 4], [1, 2, 2, 1, 1], 'float32');
+    const a = tf.complex(real, imag);
+    const b = tf.onesLike(a);
+    expect(b.dtype).toBe('complex64');
+    expect(b.shape).toEqual([1, 2, 2, 1, 1]);
+    expectArraysEqual(b, [1, 0, 1, 0, 1, 0, 1, 0]);
+  });
+
   it('6D int32 dtype', () => {
     const a = tf.tensor6d([1, 2, 3, 4], [1, 2, 2, 1, 1, 1], 'int32');
     const b = tf.onesLike(a);
@@ -679,6 +729,16 @@ describeWithFlags('onesLike', ALL_ENVS, () => {
     expectArraysClose(b, [1, 1, 1, 1]);
   });
 
+  it('6D complex dtype', () => {
+    const real = tf.tensor6d([1, 2, 3, 4], [1, 2, 2, 1, 1, 1], 'float32');
+    const imag = tf.tensor6d([1, 2, 3, 4], [1, 2, 2, 1, 1, 1], 'float32');
+    const a = tf.complex(real, imag);
+    const b = tf.onesLike(a);
+    expect(b.dtype).toBe('complex64');
+    expect(b.shape).toEqual([1, 2, 2, 1, 1, 1]);
+    expectArraysEqual(b, [1, 0, 1, 0, 1, 0, 1, 0]);
+  });
+
   it('throws when passed a non-tensor', () => {
     expect(() => tf.onesLike({} as tf.Tensor))
         .toThrowError(/Argument 'x' passed to 'onesLike' must be a Tensor/);

src/ops/tensor_ops.ts

@@ -22,7 +22,7 @@ import {TensorLike, TensorLike1D, TensorLike2D, TensorLike3D, TensorLike4D, Tens
 import {DataType, Rank, ShapeMap} from '../types';
 import {assert, assertNonNull, assertNonNegativeIntegerDimensions, flatten, inferDtype, isTypedArray, makeOnesTypedArray, makeZerosTypedArray, sizeFromShape, toTypedArray} from '../util';
 
-import {complex} from './complex_ops';
+import {complex, real, imag} from './complex_ops';
 import {op} from './operation';
 
 /**
@@ -447,6 +447,11 @@ 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');
+  if ($x.dtype === 'complex64') {
+    const r = onesLike(real($x));
+    const i = zerosLike(imag($x));
+    return complex(r, i);
+  }
   return ENV.engine.runKernel(backend => backend.onesLike($x), {$x}, null) as T;
 }