def check_results_versus_brute_force(
      self, x, axis, max_lags, center, normalize):
    """Compute auto-correlation by brute force, then compare to tf result."""
    # Brute for auto-corr -- avoiding fft and transpositions.
    axis_len = x.shape[axis]
    if max_lags is None:
      max_lags = axis_len - 1
    else:
      max_lags = min(axis_len - 1, max_lags)
    auto_corr_at_lag = []
    if center:
      x -= x.mean(axis=axis, keepdims=True)
    for m in range(max_lags + 1):
      auto_corr_at_lag.append((
          np.take(x, indices=range(0, axis_len - m), axis=axis) *
          np.conj(np.take(x, indices=range(m, axis_len), axis=axis))
      ).mean(axis=axis, keepdims=True))
    rxx = np.concatenate(auto_corr_at_lag, axis=axis)
    if normalize:
      rxx /= np.take(rxx, [0], axis=axis)

    x_ph = tf.placeholder_with_default(
        x, shape=x.shape if self.use_static_shape else None)
    with spectral_ops_test_util.fft_kernel_label_map():
      auto_corr = tfd.auto_correlation(
          x_ph,
          axis=axis,
          max_lags=max_lags,
          center=center,
          normalize=normalize)
      if self.use_static_shape:
        output_shape = list(x.shape)
        output_shape[axis] = max_lags + 1
        self.assertAllEqual(output_shape, auto_corr.shape)
      self.assertAllClose(rxx, self.evaluate(auto_corr), rtol=1e-5, atol=1e-5)

  def test_shapes(self):
    with spectral_ops_test_util.fft_kernel_label_map(), (
        self.test_session(use_gpu=True)):
      signal = np.zeros((512,)).astype(np.float32)

      # If fft_length is not provided, the smallest enclosing power of 2 of
      # frame_length (8) is used.
      stft = spectral_ops.stft(signal, frame_length=7, frame_step=8,
                               pad_end=True)
      self.assertAllEqual([64, 5], stft.shape.as_list())
      self.assertAllEqual([64, 5], stft.eval().shape)

      stft = spectral_ops.stft(signal, frame_length=8, frame_step=8,
                               pad_end=True)
      self.assertAllEqual([64, 5], stft.shape.as_list())
      self.assertAllEqual([64, 5], stft.eval().shape)

      stft = spectral_ops.stft(signal, frame_length=8, frame_step=8,
                               fft_length=16, pad_end=True)
      self.assertAllEqual([64, 9], stft.shape.as_list())
      self.assertAllEqual([64, 9], stft.eval().shape)

      stft = np.zeros((32, 9)).astype(np.complex64)

      inverse_stft = spectral_ops.inverse_stft(stft, frame_length=8,
                                               fft_length=16, frame_step=8)
      expected_length = (stft.shape[0] - 1) * 8 + 8
      self.assertAllEqual([None], inverse_stft.shape.as_list())
      self.assertAllEqual([expected_length], inverse_stft.eval().shape)

  def test_gradients_numerical(self):
    with spectral_ops_test_util.fft_kernel_label_map(), (
        self.test_session(use_gpu=True)):
      # Tuples of (signal_length, frame_length, frame_step, fft_length,
      # stft_bound, inverse_stft_bound).
      # TODO(rjryan): Investigate why STFT gradient error is so high.
      test_configs = [
          (64, 16, 8, 16),
          (64, 16, 16, 16),
          (64, 16, 7, 16),
          (64, 7, 4, 9),
          (29, 5, 1, 10),
      ]

      for (signal_length, frame_length, frame_step, fft_length) in test_configs:
        signal_shape = [signal_length]
        signal = random_ops.random_uniform(signal_shape)
        stft_shape = [max(0, 1 + (signal_length - frame_length) // frame_step),
                      fft_length // 2 + 1]
        stft = spectral_ops.stft(signal, frame_length, frame_step, fft_length,
                                 pad_end=False)
        inverse_stft_shape = [(stft_shape[0] - 1) * frame_step + frame_length]
        inverse_stft = spectral_ops.inverse_stft(stft, frame_length, frame_step,
                                                 fft_length)
        stft_error = test.compute_gradient_error(signal, [signal_length],
                                                 stft, stft_shape)
        inverse_stft_error = test.compute_gradient_error(
            stft, stft_shape, inverse_stft, inverse_stft_shape)
        self.assertLess(stft_error, 2e-3)
        self.assertLess(inverse_stft_error, 5e-4)

  def _compare(self, signal, frame_length, frame_step, fft_length):
    with spectral_ops_test_util.fft_kernel_label_map(), (
        self.test_session(use_gpu=True)) as sess:
      actual_stft = spectral_ops.stft(
          signal, frame_length, frame_step, fft_length, pad_end=False)
      signal_ph = array_ops.placeholder(dtype=dtypes.as_dtype(signal.dtype))
      actual_stft_from_ph = spectral_ops.stft(
          signal_ph, frame_length, frame_step, fft_length, pad_end=False)

      actual_inverse_stft = spectral_ops.inverse_stft(
          actual_stft, frame_length, frame_step, fft_length)

      actual_stft, actual_stft_from_ph, actual_inverse_stft = sess.run(
          [actual_stft, actual_stft_from_ph, actual_inverse_stft],
          feed_dict={signal_ph: signal})

      actual_stft_ph = array_ops.placeholder(dtype=actual_stft.dtype)
      actual_inverse_stft_from_ph = sess.run(
          spectral_ops.inverse_stft(
              actual_stft_ph, frame_length, frame_step, fft_length),
          feed_dict={actual_stft_ph: actual_stft})

      # Confirm that there is no difference in output when shape/rank is fully
      # unknown or known.
      self.assertAllClose(actual_stft, actual_stft_from_ph)
      self.assertAllClose(actual_inverse_stft, actual_inverse_stft_from_ph)

      expected_stft = SpectralOpsTest._np_stft(
          signal, fft_length, frame_step, frame_length)
      self.assertAllClose(expected_stft, actual_stft, 1e-4, 1e-4)

      expected_inverse_stft = SpectralOpsTest._np_inverse_stft(
          expected_stft, fft_length, frame_step, frame_length)
      self.assertAllClose(
          expected_inverse_stft, actual_inverse_stft, 1e-4, 1e-4)

  def testMaxLagsThresholdLessThanNeg1SameAsNone(self):
    # Setting both means we filter out items R_k from the auto-correlation
    # sequence if k > filter_beyond_lag OR k >= j where R_j < filter_threshold.

    # x_ has correlation length 10.
    iid_x_ = rng.randn(500, 1).astype(np.float32)
    x_ = (iid_x_ * np.ones((500, 10)).astype(np.float32)).reshape((5000,))
    with self.cached_session() as sess:
      with spectral_ops_test_util.fft_kernel_label_map():
        x = tf.placeholder_with_default(
            input=x_, shape=x_.shape if self.use_static_shape else None)

        ess_none_none = tfp.mcmc.effective_sample_size(
            x, filter_threshold=None, filter_beyond_lag=None)
        ess_none_200 = tfp.mcmc.effective_sample_size(
            x, filter_threshold=None, filter_beyond_lag=200)
        ess_neg2_200 = tfp.mcmc.effective_sample_size(
            x, filter_threshold=-2., filter_beyond_lag=200)
        ess_neg2_none = tfp.mcmc.effective_sample_size(
            x, filter_threshold=-2., filter_beyond_lag=None)
        ess_none_none_, ess_none_200_, ess_neg2_200_, ess_neg2_none_ = sess.run(
            [ess_none_none, ess_none_200, ess_neg2_200, ess_neg2_none])

        # filter_threshold=-2 <==> filter_threshold=None.
        self.assertAllClose(ess_none_none_, ess_neg2_none_)
        self.assertAllClose(ess_none_200_, ess_neg2_200_)

  def testMaxLagsArgsAddInAnOrManner(self):
    # Setting both means we filter out items R_k from the auto-correlation
    # sequence if k > filter_beyond_lag OR k >= j where R_j < filter_threshold.

    # x_ has correlation length 10.
    iid_x_ = rng.randn(500, 1).astype(np.float32)
    x_ = (iid_x_ * np.ones((500, 10)).astype(np.float32)).reshape((5000,))
    with self.cached_session() as sess:
      with spectral_ops_test_util.fft_kernel_label_map():
        x = tf.placeholder_with_default(
            input=x_, shape=x_.shape if self.use_static_shape else None)

        ess_1_9 = tfp.mcmc.effective_sample_size(
            x, filter_threshold=1., filter_beyond_lag=9)
        ess_1_none = tfp.mcmc.effective_sample_size(
            x, filter_threshold=1., filter_beyond_lag=None)
        ess_none_9 = tfp.mcmc.effective_sample_size(
            x, filter_threshold=1., filter_beyond_lag=9)
        ess_1_9_, ess_1_none_, ess_none_9_ = sess.run(
            [ess_1_9, ess_1_none, ess_none_9])

        # Since R_k = 1 for k < 10, and R_k < 1 for k >= 10,
        # filter_threshold = 1 <==> filter_beyond_lag = 9.
        self.assertAllClose(ess_1_9_, ess_1_none_)
        self.assertAllClose(ess_1_9_, ess_none_9_)

 def test_random(self):
   """Test randomly generated batches of data."""
   with spectral_ops_test_util.fft_kernel_label_map():
     with self.test_session(use_gpu=True):
       for shape in ([2, 20], [1], [2], [3], [10], [2, 20], [2, 3, 25]):
         signals = np.random.rand(*shape).astype(np.float32)
         for norm in (None, "ortho"):
           self._compare(signals, norm)

 def test_unknown_shape(self):
   """A test that the op runs when shape and rank are unknown."""
   with spectral_ops_test_util.fft_kernel_label_map():
     with self.session(use_gpu=True):
       signal = array_ops.placeholder_with_default(
           random_ops.random_normal((2, 3, 5)), tensor_shape.TensorShape(None))
       self.assertIsNone(signal.shape.ndims)
       mfcc_ops.mfccs_from_log_mel_spectrograms(signal).eval()

 def testEmpty(self):
   with spectral_ops_test_util.fft_kernel_label_map():
     for rank in VALID_FFT_RANKS:
       for dims in xrange(rank, rank + 3):
         x = np.zeros((0,) * dims).astype(np.float32)
         self.assertEqual(x.shape, self._tfFFT(x, rank).shape)
         x = np.zeros((0,) * dims).astype(np.complex64)
         self.assertEqual(x.shape, self._tfIFFT(x, rank).shape)

 def testBasic(self):
   with spectral_ops_test_util.fft_kernel_label_map():
     for np_type, tol in ((np.complex64, 1e-4), (np.complex128, 1e-8)):
       for rank in VALID_FFT_RANKS:
         for dims in xrange(rank, rank + 3):
           self._compare(
               np.mod(np.arange(np.power(4, dims)), 10).reshape(
                   (4,) * dims).astype(np_type), rank, rtol=tol, atol=tol)

 def test_random(self, shape):
   """Test randomly generated batches of data."""
   with spectral_ops_test_util.fft_kernel_label_map():
     with self.session(use_gpu=True):
       signals = np.random.rand(*shape).astype(np.float32)
       # Normalization not implemented for orthonormal.
       self._compare(signals, norm=None, dct_type=1)
       for norm in (None, "ortho"):
         self._compare(signals, norm, 2)
         self._compare(signals, norm, 3)

 def testIidRank2NormalHasFullEssMaxLagThresholdZero(self):
   # See similar test for Rank1Normal for reasoning.
   with self.test_session() as sess:
     with spectral_ops_test_util.fft_kernel_label_map():
       self._check_versus_expected_effective_sample_size(
           x_=rng.randn(5000, 2).astype(np.float32),
           expected_ess=5000,
           sess=sess,
           max_lags_threshold=0.,
           rtol=0.1)

  def testError(self):
    with spectral_ops_test_util.fft_kernel_label_map():
      for rank in VALID_FFT_RANKS:
        for dims in xrange(0, rank):
          x = np.zeros((1,) * dims).astype(np.complex64)
          with self.assertRaisesWithPredicateMatch(
              ValueError, "Shape .* must have rank at least {}".format(rank)):
            self._tfFFT(x, rank)
          with self.assertRaisesWithPredicateMatch(
              ValueError, "Shape .* must have rank at least {}".format(rank)):
            self._tfIFFT(x, rank)
        for dims in xrange(rank, rank + 2):
          x = np.zeros((1,) * rank)

          # Test non-rank-1 fft_length produces an error.
          fft_length = np.zeros((1, 1)).astype(np.int32)
          with self.assertRaisesWithPredicateMatch(ValueError,
                                                   "Shape .* must have rank 1"):
            self._tfFFT(x, rank, fft_length)
          with self.assertRaisesWithPredicateMatch(ValueError,
                                                   "Shape .* must have rank 1"):
            self._tfIFFT(x, rank, fft_length)

          # Test wrong fft_length length.
          fft_length = np.zeros((rank + 1,)).astype(np.int32)
          with self.assertRaisesWithPredicateMatch(
              ValueError, "Dimension must be .*but is {}.*".format(rank + 1)):
            self._tfFFT(x, rank, fft_length)
          with self.assertRaisesWithPredicateMatch(
              ValueError, "Dimension must be .*but is {}.*".format(rank + 1)):
            self._tfIFFT(x, rank, fft_length)

        # Test that calling the kernel directly without padding to fft_length
        # produces an error.
        rffts_for_rank = {
            1: [gen_spectral_ops.rfft, gen_spectral_ops.irfft],
            2: [gen_spectral_ops.rfft2d, gen_spectral_ops.irfft2d],
            3: [gen_spectral_ops.rfft3d, gen_spectral_ops.irfft3d]
        }
        rfft_fn, irfft_fn = rffts_for_rank[rank]
        with self.assertRaisesWithPredicateMatch(
            errors.InvalidArgumentError,
            "Input dimension .* must have length of at least 6 but got: 5"):
          x = np.zeros((5,) * rank).astype(np.float32)
          fft_length = [6] * rank
          with self.cached_session():
            rfft_fn(x, fft_length).eval()

        with self.assertRaisesWithPredicateMatch(
            errors.InvalidArgumentError,
            "Input dimension .* must have length of at least .* but got: 3"):
          x = np.zeros((3,) * rank).astype(np.complex64)
          fft_length = [6] * rank
          with self.cached_session():
            irfft_fn(x, fft_length).eval()

 def testGrad_Random(self):
   with spectral_ops_test_util.fft_kernel_label_map():
     for np_type, tol in ((np.float32, 1e-2), (np.float64, 1e-10)):
       for rank in VALID_FFT_RANKS:
         for dims in xrange(rank, rank + 2):
           re = np.random.rand(*((3,) * dims)).astype(np_type) * 2 - 1
           im = np.random.rand(*((3,) * dims)).astype(np_type) * 2 - 1
           self._checkGradComplex(self._tfFFTForRank(rank), re, im,
                                  rtol=tol, atol=tol)
           self._checkGradComplex(self._tfIFFTForRank(rank), re, im,
                                  rtol=tol, atol=tol)

 def testIidRank2NormalHasFullEssMaxLags10(self):
   # See similar test for Rank1Normal for reasoning.
   with self.cached_session() as sess:
     with spectral_ops_test_util.fft_kernel_label_map():
       self._check_versus_expected_effective_sample_size(
           x_=rng.randn(5000, 2).astype(np.float32),
           expected_ess=5000,
           sess=sess,
           filter_beyond_lag=10,
           filter_threshold=None,
           rtol=0.3)

 def testGrad_Simple(self):
   with spectral_ops_test_util.fft_kernel_label_map():
     for np_type, tol in ((np.float32, 1e-4), (np.float64, 1e-10)):
       for rank in VALID_FFT_RANKS:
         for dims in xrange(rank, rank + 2):
           re = np.ones(shape=(4,) * dims, dtype=np_type) / 10.0
           im = np.zeros(shape=(4,) * dims, dtype=np_type)
           self._checkGradComplex(self._tfFFTForRank(rank), re, im,
                                  rtol=tol, atol=tol)
           self._checkGradComplex(self._tfIFFTForRank(rank), re, im,
                                  rtol=tol, atol=tol)

 def test_constant_sequence_axis_0_max_lags_none_center_true(self):
   x_ = np.array([[0., 0., 0.], [1., 1., 1.]]).astype(self.dtype)
   x_ph = tf.placeholder_with_default(
       input=x_, shape=x_.shape if self.use_static_shape else None)
   with spectral_ops_test_util.fft_kernel_label_map():
     # Setting normalize = True means we divide by zero.
     auto_corr = tfp.stats.auto_correlation(
         x_ph, axis=1, normalize=False, center=True)
     if self.use_static_shape:
       self.assertEqual((2, 3), auto_corr.shape)
     auto_corr_ = self.evaluate(auto_corr)
     self.assertAllClose([[0., 0., 0.], [0., 0., 0.]], auto_corr_)

  def testRandom(self):
    with spectral_ops_test_util.fft_kernel_label_map():
      for np_type, tol in ((np.complex64, 1e-4), (np.complex128, 5e-6)):
        def gen(shape):
          n = np.prod(shape)
          re = np.random.uniform(size=n)
          im = np.random.uniform(size=n)
          return (re + im * 1j).reshape(shape)

        for rank in VALID_FFT_RANKS:
          for dims in xrange(rank, rank + 3):
            self._compare(gen((4,) * dims).astype(np_type), rank,
                          rtol=tol, atol=tol)

  def test_stft_round_trip(self):
    # Tuples of (signal_length, frame_length, frame_step, fft_length,
    # threshold, corrected_threshold).
    test_configs = [
        # 87.5% overlap.
        (4096, 256, 32, 256, 1e-5, 1e-6),
        # 75% overlap.
        (4096, 256, 64, 256, 1e-5, 1e-6),
        # Odd frame hop.
        (4096, 128, 25, 128, 1e-3, 1e-6),
        # Odd frame length.
        (4096, 127, 32, 128, 1e-3, 1e-6),
        # 50% overlap.
        (4096, 128, 64, 128, 0.40, 1e-6),
    ]

    for (signal_length, frame_length, frame_step, fft_length, threshold,
         corrected_threshold) in test_configs:
      # Generate a random white Gaussian signal.
      signal = random_ops.random_normal([signal_length])

      with spectral_ops_test_util.fft_kernel_label_map(), (
          self.test_session(use_gpu=True)) as sess:
        stft = spectral_ops.stft(signal, frame_length, frame_step, fft_length,
                                 pad_end=False)
        inverse_stft = spectral_ops.inverse_stft(stft, frame_length, frame_step,
                                                 fft_length)
        inverse_stft_corrected = spectral_ops.inverse_stft(
            stft, frame_length, frame_step, fft_length,
            window_fn=spectral_ops.inverse_stft_window_fn(frame_step))
        signal, inverse_stft, inverse_stft_corrected = sess.run(
            [signal, inverse_stft, inverse_stft_corrected])

        # Truncate signal to the size of inverse stft.
        signal = signal[:inverse_stft.shape[0]]

        # Ignore the frame_length samples at either edge.
        signal = signal[frame_length:-frame_length]
        inverse_stft = inverse_stft[frame_length:-frame_length]
        inverse_stft_corrected = inverse_stft_corrected[
            frame_length:-frame_length]

        # Check that the inverse and original signal are close up to a scale
        # factor.
        inverse_stft_scaled = inverse_stft / np.mean(np.abs(inverse_stft))
        signal_scaled = signal / np.mean(np.abs(signal))
        self.assertLess(np.std(inverse_stft_scaled - signal_scaled), threshold)

        # Check that the inverse with correction and original signal are close.
        self.assertLess(np.std(inverse_stft_corrected - signal),
                        corrected_threshold)

 def testIidRank1NormalHasFullEssMaxLags10(self):
   # With a length 5000 iid normal sequence, and max_lags = 10, we should
   # have a good estimate of ESS, and it should be close to the full sequence
   # length of 5000.
   # The choice of max_lags = 10 is a short cutoff, reasonable only since we
   # know the correlation length should be zero right away.
   with self.test_session() as sess:
     with spectral_ops_test_util.fft_kernel_label_map():
       self._check_versus_expected_effective_sample_size(
           x_=rng.randn(5000).astype(np.float32),
           expected_ess=5000,
           sess=sess,
           max_lags=10,
           rtol=0.3)