Add IRFFT ops. (tensorflow#1395)

FEATURE

Author: Lewuathe

src/ops/spectral_ops.ts

@@ -20,6 +20,7 @@ import {complex, imag, real} from '../ops/complex_ops';
 import {op} from '../ops/operation';
 import {Tensor} from '../tensor';
 import {assert} from '../util';
+import {scalar} from './tensor_ops';
 
 /**
  * Fast Fourier transform.
@@ -133,6 +134,55 @@ function rfft_(input: Tensor): Tensor {
       .reshape(outputShape);
 }
 
+/**
+ * Inversed real value input fast Fourier transform.
+ *
+ * Computes the 1-dimensional inversed discrete Fourier transform over the
+ * inner-most dimension of the real input.
+ *
+ * ```js
+ * const real = tf.tensor1d([1, 2, 3]);
+ *
+ * x.irfft().print();
+ * ```
+ * @param input The real value input to compute an irfft over.
+ */
+/**
+ * @doc {heading: 'Operations', subheading: 'Spectral', namespace: 'spectral'}
+ */
+function irfft_(input: Tensor): Tensor {
+  const innerDimensionSize = input.shape[input.shape.length - 1];
+  const batch = input.size / innerDimensionSize;
+
+  if (innerDimensionSize <= 2) {
+    const complexInput = input.as2D(batch, innerDimensionSize);
+    const ret = ENV.engine.runKernel(
+        backend => backend.ifft(complexInput), {complexInput});
+    return real(ret);
+  } else {
+    // The length of unique components of the DFT of a real-valued signal
+    // is 2 * (input_len - 1)
+    const outputShape = [batch, 2 * (innerDimensionSize - 1)];
+    const realInput = real(input).as2D(batch, innerDimensionSize);
+    const imagInput = imag(input).as2D(batch, innerDimensionSize);
+
+    const realConjugate =
+        realInput.slice([0, 1], [batch, innerDimensionSize - 2]).reverse(1);
+    const imagConjugate =
+        imagInput.slice([0, 1], [batch, innerDimensionSize - 2])
+            .reverse(1)
+            .mul(scalar(-1));
+
+    const r = realInput.concat(realConjugate, 1);
+    const i = imagInput.concat(imagConjugate, 1);
+    const complexInput = complex(r, i).as2D(outputShape[0], outputShape[1]);
+    const ret = ENV.engine.runKernel(
+        backend => backend.ifft(complexInput), {complexInput});
+    return real(ret);
+  }
+}
+
 export const fft = op({fft_});
 export const ifft = op({ifft_});
 export const rfft = op({rfft_});
+export const irfft = op({irfft_});

src/ops/spectral_ops_test.ts

@@ -214,8 +214,93 @@ describeWithFlags('2D RFFT', WEBGL_ENVS, () => {
   });
 
   it('should return the same value with TensorFlow (2x2x2 elements)', () => {
-    const t1Real = tf.tensor3d([ 1, 2, 3, 4, 5, 6, 7, 8 ], [ 2, 2, 2 ]);
-    expectArraysClose(tf.spectral.rfft(t1Real),
-                      [ 3, 0, -1, 0, 7, 0, -1, 0, 11, 0, -1, 0, 15, 0, -1, 0 ]);
+    const t1Real = tf.tensor3d([1, 2, 3, 4, 5, 6, 7, 8], [2, 2, 2]);
+    expectArraysClose(
+        tf.spectral.rfft(t1Real),
+        [3, 0, -1, 0, 7, 0, -1, 0, 11, 0, -1, 0, 15, 0, -1, 0]);
+  });
+});
+
+describeWithFlags('1D IRFFT', ALL_ENVS, () => {
+  it('should return the same value with TensorFlow (2 elements)', () => {
+    const t1Real = tf.tensor1d([1, 2]);
+    const t1Imag = tf.tensor1d([0, 0]);
+    const t1 = tf.complex(t1Real, t1Imag);
+    expectArraysClose(tf.spectral.irfft(t1), [1.5, -0.5]);
+  });
+
+  it('should return the same value with TensorFlow (5 elements)', () => {
+    const t1Real = tf.tensor1d([1, 2, 3, 4, 5]);
+    const t1Imag = tf.tensor1d([0, 0, 0, 0, 0]);
+    const t1 = tf.complex(t1Real, t1Imag);
+    expectArraysClose(
+        tf.spectral.irfft(t1),
+        [3, -0.8535534, 0, -0.14644662, 0, -0.14644662, 0, -0.8535534]);
   });
+
+  it('should return the same value with TensorFlow (5 elements) with imag',
+     () => {
+       const t1Real = tf.tensor1d([1, 2, 3, 4, 5]);
+       const t1Imag = tf.tensor1d([1, 2, 3, 4, 5]);
+       const t1 = tf.complex(t1Real, t1Imag);
+       expectArraysClose(
+           tf.spectral.irfft(t1),
+           [3, -2.6642137, 0.5, -0.45710677, 0, 0.16421354, -0.5, 0.95710677]);
+     });
+});
+
+describeWithFlags('2D IRFFT', ALL_ENVS, () => {
+  it('should return the same value with TensorFlow (2x2 elements)', () => {
+    const t1Real = tf.tensor2d([1, 2, 3, 4], [2, 2]);
+    const t1Imag = tf.tensor2d([0, 0, 0, 0], [2, 2]);
+    const t1 = tf.complex(t1Real, t1Imag);
+    expectArraysClose(tf.spectral.irfft(t1), [1.5, -0.5, 3.5, -0.5]);
+  });
+
+  it('should return the same value with TensorFlow (2x3 elements)', () => {
+    const t1Real = tf.tensor2d([1, 2, 3, 4, 5, 6], [2, 3]);
+    const t1Imag = tf.tensor2d([0, 0, 0, 0, 0, 0], [2, 3]);
+    const t1 = tf.complex(t1Real, t1Imag);
+    expectArraysClose(
+        tf.spectral.irfft(t1), [2, -0.5, 0, -0.5, 5, -0.5, 0, -0.5]);
+  });
+
+  it('should return the same value with TensorFlow (2x3 elements) with imag',
+     () => {
+       const t1Real = tf.tensor2d([1, 2, 3, 4, 5, 6], [2, 3]);
+       const t1Imag = tf.tensor2d([1, 2, 3, 4, 5, 6], [2, 3]);
+       const t1 = tf.complex(t1Real, t1Imag);
+       expectArraysClose(tf.spectral.irfft(t1), [2, -1.5, 0, 0.5, 5, -3, 0, 2]);
+     });
+});
+
+describeWithFlags('3D IRFFT', ALL_ENVS, () => {
+  it('should return the same value with TensorFlow (2x2x2 elements)', () => {
+    const t1Real = tf.tensor3d([1, 2, 3, 4, 1, 2, 3, 4], [2, 2, 2]);
+    const t1Imag = tf.tensor3d([0, 0, 0, 0, 0, 0, 0, 0], [2, 2, 2]);
+    const t1 = tf.complex(t1Real, t1Imag);
+    expectArraysClose(
+        tf.spectral.irfft(t1), [1.5, -0.5, 3.5, -0.5, 1.5, -0.5, 3.5, -0.5]);
+  });
+
+  it('should return the same value with TensorFlow (2x2x3 elements)', () => {
+    const t1Real = tf.tensor3d([1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6], [2, 2, 3]);
+    const t1Imag = tf.tensor3d([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [2, 2, 3]);
+    const t1 = tf.complex(t1Real, t1Imag);
+    expectArraysClose(tf.spectral.irfft(t1), [
+      2, -0.5, 0, -0.5, 5, -0.5, 0, -0.5, 2, -0.5, 0, -0.5, 5, -0.5, 0, -0.5
+    ]);
+  });
+
+  it('should return the same value with TensorFlow (2x2x3 elements) with imag',
+     () => {
+       const t1Real =
+           tf.tensor3d([1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6], [2, 2, 3]);
+       const t1Imag =
+           tf.tensor3d([1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6], [2, 2, 3]);
+       const t1 = tf.complex(t1Real, t1Imag);
+       expectArraysClose(
+           tf.spectral.irfft(t1),
+           [2, -1.5, 0, 0.5, 5, -3, 0, 2, 2, -1.5, 0, 0.5, 5, -3, 0, 2]);
+     });
 });